tensor¶

The core of Tricycle is the Tensor object, which is implemented in this file.

A Tensor is a wrapper around a numpy/cupy array that adds automatic differentiation.

The autodiff algorithm itself can be found in Tensor.backward.

This file also contains a few other helpful functions like batch which converts tensors to batched tensors.

class Tensor(array, requires_grad=True, is_batched=False, args=None, back_fns=None, dtype=None, name=None, _id=None)[source]¶

Bases: object

An N-dimensional grid of numbers. This is implemented as a subclass of a standard numpy array.

Parameters:

array (_SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | bool | int | float | complex | str | bytes | _NestedSequence[bool | int | float | complex | str | bytes])
requires_grad (bool)
is_batched (bool)
args (tuple[Tensor, ...] | None)
back_fns (tuple[Op, ...] | None)
dtype (dtype[Any] | None | type[Any] | _SupportsDType[dtype[Any]] | str | tuple[Any, int] | tuple[Any, SupportsIndex | Sequence[SupportsIndex]] | list[Any] | _DTypeDict | tuple[Any, Any])
name (str | None)
_id (int | None)

_id¶

Unique identifier for the tensor.

Type:: int

_parents¶

Parent tensors in the computation graph.

Type:: set[Tensor] | None

array¶

The underlying numpy/cupy array.

Type:: ArrayLike

args¶

Arguments used to create this tensor.

Type:: tuple[Tensor, …] | None

back_fns¶

Backward functions for gradient computation.

Type:: tuple[Op, …] | None

grad¶

Gradient of this tensor.

Type:: Optional[Tensor]

name¶

Name of the tensor.

Type:: Optional[str]

requires_grad¶

Whether this tensor requires gradient computation.

Type:: bool

is_batched¶

Whether this tensor is batched.

Type:: bool

backward(clip=None)[source]¶

Performs a backward pass through the graph, calculating the gradient for each parameter.

Parameters:: clip (float | None, optional) – Maximum absolute value for gradient clipping. Defaults to None.

close_to(other, equal_nan=False, rtol=0.0001, **kwargs)[source]¶

Checks if this tensor is close to another tensor or value within a tolerance.

Parameters:

other (Union[Tensor, ArrayLike, float, int]) – The tensor or value to compare against.
equal_nan (bool, optional) – Whether to consider NaN values as equal. Defaults to False.
rtol (float, optional) – The relative tolerance parameter. Defaults to 1e-4.
**kwargs – Additional keyword arguments to pass to numpy.allclose or cupy.allclose.

Returns:

True if the tensors are close, False otherwise.

Return type:

bool

property dtype: dtype¶

Returns the data type of the tensor.

Returns:: The data type of the tensor.
Return type:: np.dtype

einsum(subscript)[source]¶

Performs an einsum operation on the tensor.

Parameters:: subscript (str) – The einsum subscript string.
Returns:: The result of the einsum operation.
Return type:: Tensor

from_batched()[source]¶

Treats a batched tensor as a normal, non-batched, tensor.

Returns:: A new non-batched tensor.
Return type:: Tensor

from_gpu()[source]¶

Moves this tensor from the GPU to CPU.

Returns:: The tensor moved to the CPU.
Return type:: Tensor
Raises:: GPUDisabledException – If CuPY is not enabled.

mean()[source]¶

Computes the mean of all elements in the tensor.

Returns:: A new tensor containing the mean value.
Return type:: Tensor

property ndim: int¶

Returns the number of dimensions of the tensor.

Returns:: The number of dimensions.
Return type:: int

numpy()[source]¶

Returns the underlying array as a numpy array.

Returns:: The tensor data as a numpy array.
Return type:: np.ndarray

property on_gpu¶

Checks if the tensor is currently on the GPU.

Returns:: True if the tensor is on the GPU, False otherwise.
Return type:: bool

repeat(n_repeats)[source]¶

Repeats the tensor.

Parameters:: n_repeats (int) – The number of times to repeat the tensor.
Returns:: The repeated tensor.
Return type:: Tensor

reshape(shape)[source]¶

Reshapes the tensor to the given shape.

Parameters:: shape (Sequence[int]) – The new shape for the tensor.
Returns:: The reshaped tensor.
Return type:: Tensor

property shape: Sequence[int]¶

Returns the shape of the tensor.

Returns:: The shape of the tensor.
Return type:: Sequence[int]

split(n_splits, axis=-1)[source]¶

Splits the tensor into multiple sub-tensors.

Parameters:

n_splits (int) – The number of splits to perform.
axis (int, optional) – The axis along which to split. Defaults to -1.

Returns:

A list of split tensors.

Return type:

List[Tensor]

sum()[source]¶

Computes the sum of all elements in the tensor.

Returns:: A new tensor containing the sum.
Return type:: Tensor

to_batched()[source]¶

Treats this tensor as a batch of tensors.

Returns:: A new batched tensor.
Return type:: Tensor

to_gpu(device=0)[source]¶

Moves this tensor to the GPU, if cupy is enabled.

Parameters:: device (int, optional) – The GPU device number. Defaults to 0.
Returns:: The tensor moved to the GPU.
Return type:: Tensor
Raises:: GPUDisabledException – If CuPY is not enabled.

property xp¶

Returns the appropriate array library (numpy or cupy) for the tensor.

Returns:: The array library (numpy or cupy).
Return type:: module

zero_grad()[source]¶

Removes any gradients or references to other tensors.

Returns:: The tensor with gradients and references cleared.
Return type:: Tensor

select_backend(*tensors)[source]¶

Given some tensors, if any of them are on the GPU, return the cupy backend. Otherwise default to the numpy backend.

Parameters:: *tensors (Tensor | ndarray | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | bool | int | float | complex | str | bytes | _NestedSequence[bool | int | float | complex | str | bytes]) – Variable number of tensors or arrays to check.
Returns:: The appropriate backend module (numpy or cupy).
Return type:: module