tensor

Utility functions for PyTorch tensors.

Functions

numpy_to_torch

Convert numpy arrays to torch tensors.

to_empty_if_meta_device

Move tensors to device if not meta device; otherwise materialize with empty_like().

torch_detach

Try to recursively detach the data from the computation graph.

torch_to

Try to recursively move the data to the specified args/kwargs.

torch_to_numpy

Convert torch tensors to numpy arrays.

numpy_to_torch(np_outputs)

Convert numpy arrays to torch tensors.

Parameters:

np_outputs (list[ndarray])

Return type:

list[Tensor]

to_empty_if_meta_device(module, *, device, recurse=True)

Move tensors to device if not meta device; otherwise materialize with empty_like().

Officially, torch suggests to_empty() for meta device materialization. Under the hood, torch.empty_like() is applied to all parameters or buffers (see _apply). This may accidently overwrite buffers with precomputed values during construction. Given the goal is to only materialize those tensors on meta device, this function checks the device first and only move the tensor to the destination if it is not on meta device.

Parameters:
  • module (Module) – The target module to apply this transformation.

  • device (device) – The desired device of the parameters and buffers in this module.

  • recurse – Whether parameters and buffers of submodules should be recursively moved to the specified device.

torch_detach(data)

Try to recursively detach the data from the computation graph.

torch_to(data, *args, **kwargs)

Try to recursively move the data to the specified args/kwargs.

torch_to_numpy(inputs)

Convert torch tensors to numpy arrays.

Parameters:

inputs (list[Tensor])

Return type:

list[ndarray]