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]