blazefl.core.process_tensors_in_object

blazefl.core.process_tensors_in_object(obj: T, mode: Literal['move', 'replace'], max_depth: int = 10) → T

Recursively traverses an object to process torch.Tensor instances.

This function handles different data structures, including dictionaries, lists, tuples, and general objects with a __dict__ attribute. It also manages circular references to prevent infinite recursion.

Parameters:

• obj – The object to process.

• mode –

  • “move”: Moves tensors to shared memory in-place by calling

    .share_memory_(). This is suitable for sending data from a parent to a worker process.

  • ”replace”: Creates a deep copy of the object, replacing all

    tensors with a lightweight SHMHandle(). This is used to create a serializable “receipt” object that can be returned from a worker to the parent without copying tensor data.

• max_depth – The maximum recursion depth. Defaults to 10.

Returns:

• In “move” mode, returns the original object, modified in-place.

• In “replace” mode, returns a new object where tensors have been replaced with SHMHandle instances.