tinygrad shape.shapetracker

Note

You likely want the upstream tinygrad, not tinygrab. Tinygrab contains AI generated docstrings for a tinygrad snapshot. Upstream: https://tinygrad.org

class tinygrad.shape.shapetracker.ShapeTracker(views: Tuple[View, ...])[source]

Bases: object

This class is responsible for tracking the shape of a View object. It contains methods to create instances from shapes, check if instances are contiguous and obtain various properties like shape, size and variables.

views

A tuple of View objects representing the shape.

Type:: Tuple[View, …]

axis_is_masked(axis: int) → bool[source]

This method checks if an axis is masked.

Parameters:: axis (int) – The axis to check for masking.
Returns:: True if the axis is masked, False otherwise.
Return type:: bool

property contiguous: bool

Check if the current ShapeTracker is contiguous, meaning that it contains exactly one View and that View is contiguous.

Returns:: True if the ShapeTracker is contiguous, False otherwise.
Return type:: bool

expand(new_shape: Tuple[Node | int, ...]) → ShapeTracker[source]

This method expands the shape tracker to a new shape.

Parameters:: new_shape (Tuple[sint, ...]) – The new shape for the shape tracker.
Returns:: The expanded shape tracker.
Return type:: ShapeTracker

expr_idxs(idxs: Iterable[Node] | None = None)[source]

Calculate the expression indices and validity.

self: The object itself

idxs

An optional iterable of nodes

Type:: Optional[Iterable[Node]]

expr_node(idx: Node | str = 'idx')[source]

This method creates an expression node.

Parameters:: idx (Union[Node, str]) – The index to use for the expression node. Default is “idx”.
Returns:: The created expression node.
Return type:: Node

static from_shape(shape: Tuple[Node | int, ...])[source]

Create a ShapeTracker instance from a given shape.

Parameters:: shape (Tuple[sint, ...]) – The shape to create the ShapeTracker for.
Returns:: A new ShapeTracker instance.
Return type:: ShapeTracker

pad(arg: Tuple[Tuple[int, int], ...]) → ShapeTracker[source]

This method pads the shape tracker with the given arguments.

Parameters:: arg (Tuple[Tuple[int, int], ...]) – The padding argument.
Returns:: The padded shape tracker.
Return type:: ShapeTracker

permute(axis: Tuple[int, ...]) → ShapeTracker[source]

This method permutes the shape tracker according to a given axis order.

Parameters:: axis (Tuple[int, ...]) – The new axis order for the shape tracker.
Returns:: The permuted shape tracker.
Return type:: ShapeTracker

real_strides(ignore_valid=False) → Tuple[Node | int | None, ...][source]

Calculate the strides of the shape tracker.

If a stride is not always valid, it will be None.

Arributes:: self: The object itself ignore_valid (bool): Whether to ignore validity or not

Returns:: A tuple of strides for each dimension
Return type:: Tuple[Optional[sint], …]

reshape(new_shape: Tuple[Node | int, ...]) → ShapeTracker[source]

This method reshapes the shape tracker to a new shape.

Parameters:: new_shape (Tuple[sint, ...]) – The new shape for the shape tracker.
Returns:: The reshaped shape tracker.
Return type:: ShapeTracker

property shape: Tuple[Node | int, ...]

Get the shape of the last View in the ShapeTracker.

Returns:: The shape of the last View.
Return type:: Tuple[sint, …]

shrink(arg: Tuple[Tuple[Node | int, Node | int], ...]) → ShapeTracker[source]

This method shrinks the shape tracker with the given arguments.

Parameters:: arg (Tuple[Tuple[sint, sint], ...]) – The shrinking argument.
Returns:: The shrunk shape tracker.
Return type:: ShapeTracker

simplify() → ShapeTracker[source]

Simplify the shape tracker.

self: The object itself

Returns:: The simplified shape tracker
Return type:: ShapeTracker

size() → int[source]

Calculate the size of the ShapeTracker by finding the maximum index in the expression and adding one.

Returns:: The size of the ShapeTracker.
Return type:: int

stride(mul: Tuple[int, ...]) → ShapeTracker[source]

This method applies striding to the shape tracker with a given multiplier.

Parameters:: mul (Tuple[int, ...]) – The stride multiplier.
Returns:: The strided shape tracker.
Return type:: ShapeTracker

to_movement_ops() → List[Tuple[MovementOps, Tuple]][source]

Convert the shape tracker object into a list of movement operations and their corresponding arguments.

Returns:: A list where each element is a tuple containing a movement operation and its corresponding arguments.
Return type:: List[Tuple[MovementOps, Tuple]]

unbind() → ShapeTracker[source]

Unbind the shape tracker object into its constituent parts.

Returns:: A new shape tracker object with the same properties as the original, but unbound.
Return type:: ShapeTracker

unit_stride_axes(ignore_valid=False) → List[int][source]

Return a list of axes with unit strides.

self: The object itself

ignore_valid

Whether to ignore validity or not

Type:: bool

Returns:: A list of axes with unit strides
Return type:: List[int]

property var_vals: Dict[Variable, int]

Return a dictionary mapping variables to their integer values.

Returns:: A dictionary where each key-value pair represents a variable and its corresponding integer value.
Return type:: Dict[Variable, int]

vars() → Set[Variable][source]

Get the union of all variables in each View within the ShapeTracker.

Returns:: The set of unique variables in the ShapeTracker.
Return type:: Set[Variable]

views: Tuple[View, ...]

tinygrad.shape.shapetracker.expr_idxs(view: View, idxs: Tuple[Node, ...]) → Node[source]

Generate an expression for a given view and list of indices.

Parameters:

view (View) – The input view.
idxs (Tuple[Node, ...]) – Indices corresponding to each dimension in the view’s shape.

Returns:

The generated expression.

Return type:

Node

Raises:

AssertionError – If the number of indices does not match the number of dimensions in the view.

tinygrad.shape.shapetracker.expr_node(view: View, idx: Node | None = None) → Node[source]

Generate an expression for a given view and index node.

tinygrad.shape.shapetracker.view: The view to generate the expression for.

tinygrad.shape.shapetracker.idx: An optional index node. If not provided, a new one is created with the range of 0 to prod(view.shape) - 1.

Returns:: A node representing the generated expression.

tinygrad.shape.shapetracker.expr_node_mask(view: View, idx: Node, valid: Node | None = None) → Node[source]

Generate an expression for a view with mask.

tinygrad.shape.shapetracker.view: The view to generate the expression for.

tinygrad.shape.shapetracker.idx: The index node.

tinygrad.shape.shapetracker.valid: An optional node representing validity.

Returns:: A node representing the generated expression.

tinygrad.shape.shapetracker.get_contraction(old_shape: Tuple[Node | int, ...], new_shape: Tuple[Node | int, ...]) → List[List[int]] | None[source]

Returns the axes to create new_shape if new_shape can be created by combining axis from old_shape.

Parameters:

old_shape (Tuple[sint, ...]) – A tuple of integers representing the original shape.
new_shape (Tuple[sint, ...]) – A tuple of integers representing the desired new shape.

Returns:

A list of lists containing the axes to create new_shape from old_shape, or None if it’s not possible.

Return type:

Optional[List[List[int]]]

Attributes:

acc_old (list): List of accumulated multiplication results of old_shape elements.
acc_new (list): List of accumulated multiplication results of new_shape elements.
split (list): List of indices where the axes to combine from old_shape are located in new_shape.

Raises:

ValueError – If it’s not possible to create new_shape from old_shape.

Example:

>>> old_shape = (2, 3, 4)
>>> new_shape = (6, 4)
>>> get_contraction(old_shape, new_shape)
[[0], [1]]

tinygrad.shape.shapetracker.idxs_to_idx(shape: Tuple[int, ...], idxs: Tuple[Node, ...]) → Node[source]

Convert a list of indices for each dimension into a single index.

Parameters:

shape (Tuple[int, ...]) – The shape of the array.
idxs (Tuple[Node, ...]) – Indices corresponding to each dimension in the shape.

Returns:

The generated expression for the single index.

Return type:

Node

Raises:

AssertionError – If the number of indices does not match the number of dimensions in the shape.

tinygrad.shape.shapetracker.merge_views(vm2: View, vm1: View) → View | None[source]

Merge two views into a single view if possible.

Parameters:

vm2 (View) – The first view to merge.
vm1 (View) – The second view to merge.

Returns:

The merged view, or None if the views could not be merged.

Return type:

Optional[View]