View & Shape#

reshape#

def reshape(x: 'Tensor', shape: 'tuple[int, ...]') -> 'Tensor':

transpose#

def transpose(x: 'Tensor', axis1: 'int', axis2: 'int') -> 'Tensor':

permute#

def permute(x: 'Tensor', order: 'tuple[int, ...]') -> 'Tensor':

unsqueeze#

def unsqueeze(x: 'Tensor', axis: 'int' = 0) -> 'Tensor':

squeeze#

def squeeze(x: 'Tensor', axis: 'int' = 0) -> 'Tensor':

flatten#

def flatten(x: 'Tensor', start_dim: 'int' = 0, end_dim: 'int' = -1) -> 'Tensor':

Flatten a range of dimensions into a single dimension using reshape.


broadcast_to#

def broadcast_to(x: 'Tensor', shape: 'tuple[int, ...]') -> 'Tensor':

concatenate#

def concatenate(tensors: 'Sequence[Tensor]', axis: 'int' = 0) -> 'Tensor':

stack#

def stack(tensors: 'list[Tensor]', axis: 'int' = 0) -> 'Tensor':

Stack tensors along a new axis.


gather#

def gather(x: 'Tensor', indices: 'Tensor', axis: 'int' = 0) -> 'Tensor':

Gather elements from x along axis using indices.


scatter#

def scatter(x: 'Tensor', indices: 'Tensor', updates: 'Tensor', axis: 'int' = 0) -> 'Tensor':

Scatter updates into x at indices along axis.


slice_tensor#

def slice_tensor(x: 'Tensor', start: 'Any', size: 'Any') -> 'Tensor':

slice_update#

def slice_update(x: 'Tensor', update: 'Tensor', start: 'Any', size: 'Any') -> 'Tensor':

Update a slice of x with new values.


moveaxis#

def moveaxis(x: 'Tensor', source: 'int', destination: 'int') -> 'Tensor':

swap_axes#

def swap_axes(x: 'Tensor', axis1: 'int', axis2: 'int') -> 'Tensor':

flip#

def flip(x: 'Tensor', axis: 'int') -> 'Tensor':

pad#

def pad(x: 'Tensor', paddings: 'list[tuple[int, int]]' = None, mode: 'str' = 'constant', value: 'float' = 0.0, **kwargs) -> 'Tensor':

rebind#

def rebind(x: 'Tensor', shape: 'tuple[int, ...]', **kwargs) -> 'Tensor':

as_interleaved_complex#

def as_interleaved_complex(x: 'Tensor') -> 'Tensor':

view_as_real_interleaved#

def view_as_real_interleaved(x: 'Tensor') -> 'Tensor':

broadcast_to_physical#

def broadcast_to_physical(x: 'Tensor', shape: 'tuple[int, ...]') -> 'Tensor':

squeeze_physical#

def squeeze_physical(x: 'Tensor', axis: 'int' = 0) -> 'Tensor':

unsqueeze_physical#

def unsqueeze_physical(x: 'Tensor', axis: 'int' = 0) -> 'Tensor':

incr_batch_dims#

def incr_batch_dims(x: 'Tensor') -> 'Tensor':

Increment batch_dims counter (first physical dim becomes batch dim).


decr_batch_dims#

def decr_batch_dims(x: 'Tensor') -> 'Tensor':

Decrement batch_dims counter (first batch dim becomes logical dim).


move_axis_to_batch_dims#

def move_axis_to_batch_dims(x: 'Tensor', axis: 'int') -> 'Tensor':

Move a logical axis into the batch dimensions (3 ops: calc + moveaxis_physical + incr).


move_axis_from_batch_dims#

def move_axis_from_batch_dims(x: 'Tensor', batch_axis: 'int' = 0, logical_destination: 'int' = 0) -> 'Tensor':

Move a batch dimension to logical axis (3 ops: calc + moveaxis_physical + decr).