tensor

cast_to_device

cast_to_device(
    value: T,
    *,
    device: str | int | device | None = None,
    dtype: dtype | None = None,
) -> T

Make a deep copy of value, casting all tensors to the given device and dtype.

Adapted from: https://github.com/pytorch/pytorch/blob/49444c3e546bf240bed24a101e747422d1f8a0ee/torch/optim/optimizer.py#L209C1-L225C29.

Parameters:

Name	Type	Description	Default
value	T	The value to recursively copy and cast.	required
device	str | int | device | None	The device to cast tensors to.	None
dtype	dtype | None	The dtype to cast tensors. Only applied to floating point tensors.	None

Returns:

Type	Description
T	The copied and casted value.

collect_devices

collect_devices(
    value: Tensor
    | dict[Any, Any]
    | UserDict[Any, Any]
    | Iterable[Any]
    | Any,
) -> set[torch.device]

Collect all devices in the given value.

Parameters:

Name	Type	Description	Default
value	Tensor | dict[Any, Any] | UserDict[Any, Any] | Iterable[Any] | Any	The value to recursively collect devices from.	required

Returns:

Type	Description
set[torch.device]	The set of all devices in the given value.

collect_floating_point_dtypes

collect_floating_point_dtypes(
    value: Tensor
    | dict[Any, Any]
    | UserDict[Any, Any]
    | Iterable[Any]
    | Any,
) -> set[torch.dtype]

Collect all floating point dtypes in the given value.

Parameters:

Name	Type	Description	Default
value	Tensor | dict[Any, Any] | UserDict[Any, Any] | Iterable[Any] | Any	The value to recursively collect floating point dtypes from.	required

Returns:

Type	Description
set[torch.dtype]	The set of all floating point dtypes in the given value.

hash_tensor_data

hash_tensor_data(tensor: Tensor) -> int

Compute the hash of the tensor's data represented as a string.

Note

Since 0 and -0 have different byte representations, they will produce different hash values.

Parameters:

Name	Type	Description	Default
tensor	Tensor	The tensor whose data to hash.	required

Returns:

Type	Description
int	The hash of the tensor's data.

Changed in version 0.76.0: Moved hash_tensor to its own utility and renamed for clarity that the hash operates on the tensor data.

to_empty

to_empty(
    module: ModuleT,
    *,
    device: device | str | int | None,
    recurse: bool = True,
) -> ModuleT

Move the parameters and buffers to the specified device without copying storage if they are not already on that device.

Parameters:

Name	Type	Description	Default
module	ModuleT	The module whose parameters and buffers to (maybe) move.	required
device	device | str | int | None	The desired device of the parameters and buffers in the module. If None, the default device is used.	required
recurse	bool	Whether parameters and buffers of submodules should be recursively moved to the specified device.	True

Returns:

Type	Description
ModuleT	The (maybe) moved module.