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.
masked_mean_over_active_rows	Average values over the active positions of each row, then over the rows that have at least one active position.
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.

masked_mean_over_active_rows

masked_mean_over_active_rows(
    values: Tensor, mask: Tensor
) -> torch.Tensor

Average values over the active positions of each row, then over the rows that have at least one active position.

This is the NaN-safe replacement for torch.masked.mean(values, mask=mask, dim=-1).mean(dim=0). That idiom performs a 0/0 reduction for any row whose mask is entirely False, yielding NaN for that row and propagating NaN through the batch reduction to the scalar result. Rows with no active positions carry no signal, so they are excluded from the batch average entirely rather than being counted as zeros; a batch in which every row is empty reduces to 0.

No NaN is ever materialized: the masked-out entries are zeroed before the summation and the per-row divisor is clamped, so the backward pass stays finite and empty rows simply receive zero gradient. Passing NaN through torch.where after the fact would not be equivalent, because a zero-weighted NaN still contaminates the gradient.

Note

This weights each row equally, independent of how many positions that row has active. It is therefore not the same as a global torch.masked.mean(values, mask=mask), which weights each active position equally.

Parameters:

Name	Type	Description	Default
values	Tensor	The tensor to reduce, of shape (batch_size, sequence_length).	required
mask	Tensor	A mask of shape (batch_size, sequence_length) selecting the active positions. Non-boolean (0/1 integer or float) masks are cast to bool, so a nonzero entry counts as active.	required

Returns:

Type	Description
torch.Tensor	A scalar tensor holding the mean over active positions of active rows.

Example

values = torch.tensor([[1.0, 3.0], [5.0, 7.0]]) mask = torch.tensor([[True, True], [True, True]]) masked_mean_over_active_rows(values, mask) tensor(4.)

A fully-masked row is dropped rather than poisoning the result:

mask = torch.tensor([[True, True], [False, False]]) masked_mean_over_active_rows(values, mask) tensor(2.)

Added in version v3.62.0. NaN-safe two-stage masked mean: reduces over active positions within each row, then over the rows that have any active position. Replaces the `torch.masked.mean(..., dim=-1).mean(dim=0)` idiom, whose 0/0 reduction on a fully-masked row returns `NaN` and poisons the whole batch.

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.