ShapedTensor

class ShapedTensor(owner: Module, name: str, value: Tensor | Parameter | None, constraints: dict[int, int] | None = None, persist_data: bool = True, persist_constraints: bool = False, strict: bool = True, live: bool = False)[source]

Bases: object

Tensor attribute with constrained shape.

Some states for value are ignored for convenience. If it is None, an instance of either UninitializedBuffer or UninitializedParameter, or has no elements and only a single dimension (such as if created with torch.empty(0)).

When value is None, a registered buffer is created, otherwise a parameter will only be added if an Parameter is given. Assignment of a parameter to None is unsupported.

Parameters:
  • owner (Module) – module to which this attribute will belong.

  • name (str) – name of the attribute.

  • value (torch.Tensor | nn.Parameter | None) – tensor-like data for the attribute.

  • constraints (dict[int, int] | None, optional) – constraints given as a dictionary of dimensions to their corresponding size. Defaults to None.

  • persist_data (bool, optional) – if the data should persist across the state dictionary, only used with buffers. Defaults to True.

  • persist_constraints (bool, optional) – if the constraints should persist across the state dictionary. Defaults to False.

  • strict (bool, optional) – if each dimension must specify a unique dimension for the tensor. Defaults to True.

  • live (bool, optional) – if constraint validity should be tested on assignment. Defaults to False.

Raises:

RuntimeError – given data is neither ignored nor compatible with given constraints.

Caution

This has a finalizer which will delete the attributes added to the module when its reference count goes to zero.

LinkedAttributes

alias of ShapedTensorAttributes

property attributes: ShapedTensorAttributes

Names of the dependent attributes created.

This is a named tuple with attributes data and constraints.

Returns:

names of the created attributes in the containing module.

Return type:

ShapedTensor.LinkedAttributes

compatible(tensor: Tensor | Parameter) bool[source]

Checks if a tensor is compatible.

Parameters:

tensor (torch.Tensor | nn.Parameter) – tensor to test for compatibility.

Returns:

if the given tensor is dimensionally compatible.

Return type:

bool

property constraints: dict[int, int]

Dictionary of registered constraints.

Each key corresponds to a dimension of the tensor, and its associated value is the size of that dimension.

Returns:

dictionary of constraints.

Return type:

dict[int, int]

classmethod create(owner: Module, name: str, value: Tensor | Parameter | None, constraints: dict[int, int] | None = None, persist_data: bool = True, persist_constraints: bool = False, strict: bool = True, live: bool = False) None[source]

Creates a shaped tensor and adds it as an attribute.

The following two calls are equivalent.

module.name = ShapedTensor(module, name, value)
ShapedTensor.create(module, name, value)
Parameters:
  • owner (Module) – module to which this attribute will belong.

  • name (str) – name of the attribute.

  • value (torch.Tensor | nn.Parameter | None) – tensor-like data for the attribute.

  • constraints (dict[int, int] | None, optional) – constraints given as a dictionary of dimensions to their corresponding size. Defaults to None.

  • persist_data (bool, optional) – if the data should persist across the state dictionary, only used with buffers. Defaults to True.

  • persist_constraints (bool, optional) – if the constraints should persist across the state dictionary. Defaults to False.

  • strict (bool, optional) – if each dimension must specify a unique dimension for the tensor. Defaults to True.

  • live (bool, optional) – if constraint validity should be tested on assignment. Defaults to False.

property dimensionality: int

Minimum number of dimensions a tensor needs to satisfy constraints.

Returns:

minimum valid dimensionality.

Return type:

int

property ignored: bool

If the current tensor is ignored.

Returns:

if the current tensor is ignored.

Return type:

bool

property live: bool

If assignments should be constraint tested.

Parameters:

value (bool) – if assignments should be constraint tested.

Returns:

if assignments should be constraint tested.

Return type:

bool

property name: str

Name of the attribute.

Two attributes with names derived from name are added to the owner.

  • _{name}_data, the constrained tensor.

  • _{name}_constraints, the dictionary of constraints.

Returns:

name of the attribute.

Return type:

str

property owner: Module | None

Module which owns this attribute.

Returns:

owner of the attribute if it exists.

Return type:

Module | None

reconstrain(dim: int, size: int | None) Tensor | Parameter | None[source]

Add, edit, or remove a constraint.

When size is None, the corresponding constraint will be removed. When dim is not in the current constraints, it will add that constraint. When dim is in the current constraints and size is not None, that constraint will be altered. Automatic resizing only occurs on editing a constraint, not on adding a constraint.

When editing a constraint, if a tensor is already compatible with that constraint, then it will not be altered. If the size of the dimension is reduced, then elements will be cut off, preserving those towards the end. If the size of the dimension is increased, then zero-valued elements will be added to the start.

Parameters:
  • dim (int) – dimension on which to modify the constraint.

  • size (int | None) – new size for the specified dimension.

Raises:
  • ValueError – dimension specified for constraint removal is unconstrained.

  • ValueError – constrained tensor would be invalidated by adding this constraint.

  • RuntimeError – previous operation invalidated constrained tensor.

  • RuntimeError – constrained tensor cannot be made valid with constraint.

Returns:

newly constrained value.

Return type:

torch.Tensor | nn.Parameter | None

static resize(value: Tensor | Parameter, dim: int, size: int, preserve_tail: bool = True, fill: Any = 0) Tensor | Parameter[source]

Resizes a tensor’s dimension.

A more generalized version of the built-in automatic resizing, this can be used before calling reconstrain() if more control is needed.

If value is a tensor, then a tensor will be returned, a new tensor if the shape needed to be changed, otherwise the same tensor as was given.

If value is a parameter, then a parameter will be returned. If it needed to be reshaped, the parameter’s data will be set with the reshaped data prior to being returned.

When preserve_tail is True, the tail of the tensor will be kept as-is, otherwise the head of the tensor will be kept. This corresponds to where slices will be removed or appended. Assuming a dimension is not sized to zero, then if preserve_tail is True, tensor[-1] will return the same values before and after, otherwise tensor[0] will.

Parameters:
  • value (torch.Tensor | nn.Parameter) – tensor-like value to resize.

  • dim (int) – dimension to resize.

  • size (int) – new size for the dimension.

  • preserve_tail (bool, optional) – if the tail (higher indices) of a tensor should be unaltered. Defaults to True.

  • fill (Any, optional) – value with which to fill expanded dimensions. Defaults to 0.

Returns:

resized tensor.

Return type:

torch.Tensor | nn.Parameter

property strict: bool

If constraints should be strictly tested.

When strict constraints are used, each constrained dimension must refer to a unique tensor dimension. For example, if constraints are placed on dimensions 0 and -1, then strict would require a tensor to have at least two dimensions. Otherwise, as long as the constraints are all met, regardless of uniqueness, a tensor is considered valid.

Parameters:

value (bool) – if constraints should be strictly tested.

Returns:

if constraints should be strictly tested.

Return type:

bool

property valid: bool

If the shaped tensor is valid.

A shaped tensor is considered valid if the value is either ignored or is compatible with the current constraints and its owner still exists.

Returns:

if the shaped tensor is valid.

Return type:

bool

property value: Tensor | Parameter | None

Value of the constrained tensor.

If live was set on initialization, every setter call will ensure the tensor being set is valid (constrained or ignored).

When created as a Parameter, assignment to None is prevented. If the current value is a Parameter but the assigned value is a Tensor, it will automatically assign to the data attribute of value.

Parameters:

(value (value) – torch.Tensor | nn.Parameter | None): value to which the constrained attribute will be set.

Returns:

constrained attribute.

Return type:

torch.Tensor | nn.Parameter | None