tensor

Module for PyTorch tensor utilities.

Functions:

Name	Description
cast_to_device	Make a deep copy of value, casting all tensors to the given device and dtype.
collect_devices	Collect all devices in the given value.
collect_floating_point_dtypes	Collect all floating point dtypes in the given value.
hash_tensor_data	Compute the hash of the tensor's data represented as a string.
prepare_for_safe_division	Prepare a tensor for safe division by adding a small value to it which is relative to the scale of the input tensor itself.
split_2d	Split a (batch_size, sequence_length, ...) tensor into chunks along the first 2 dimensions.

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.

patch_meta_data_assignment

patch_meta_data_assignment() -> Generator[None]

Intercept tensor.data = value for meta tensors and materialize them.

Swaps in a real tensor on value.device before the assignment would fail. Uses a re-entrancy reference count and a lock so that nested calls and concurrent threads behave correctly: the patch is installed on the first entry and restored only when the outermost context exits.

Warning

Scope this as narrowly as possible, e.g. around model construction/loading.

Yields:

Type	Description
Generator[None]	None.

prepare_for_safe_division

prepare_for_safe_division(
    tensor: Tensor, safety_factor: float = 100.0
) -> torch.Tensor

Prepare a tensor for safe division by adding a small value to it which is relative to the scale of the input tensor itself.

Parameters:

Name	Type	Description	Default
tensor	Tensor	The tensor to prepare for safe division.	required
safety_factor	float	A factor to multiply the small value by. This can be used to increase the value added to the tensor, providing more numerical stability at the cost of potentially more bias.	100.0

Returns:

Type	Description
torch.Tensor	The tensor with a small value added to it for safe division.

split_2d

split_2d(
    tensor: Tensor,
    max_batch_size: int | None,
    max_sequence_length: int | None,
) -> tuple[tuple[torch.Tensor, ...], ...]

Split a (batch_size, sequence_length, ...) tensor into chunks along the first 2 dimensions.

Parameters:

Name	Type	Description	Default
tensor	Tensor	The tensor to split along the first 2 dimensions.	required
max_batch_size	int | None	The max size of chunks in the first dimension. If None, no splitting is performed.	required
max_sequence_length	int | None	The max size of chunks in the second dimension. If None, no splitting is performed.	required

Returns:

Type	Description
tuple[tuple[torch.Tensor, ...], ...]	A tuple of tuple of tensors of shape (<= max_batch_size, <= max_sequence_length, ...).

Added in version v3.35.0. Porting utility function from Stained Glass Analytics to Stained Glass Core.