Mojo module

int_tuple

Hierarchical integer tuple data structures for high-performance tensor operations.

This module provides a flexible, memory-efficient implementation of nested integer tuples optimized for tensor shape, stride, and index operations in high-performance computing. The core data structures support both flat and hierarchical representations with efficient memory sharing and zero-copy views.

Key components:

IntArray: Low-level register-passable array with direct memory management
IntTuple: Hierarchical nested tuple with efficient memory layout and operations
Utility functions for tensor shape manipulation, coordinate transformations, and layout operations

Performance features:

Register-passable data structures for optimal compiler optimizations
Zero-copy views for efficient memory sharing
Specialized memory layout for nested structures
Optimized algorithms for common tensor operations

Common operations:

Shape manipulation: flatten, to_nest, apply, product, sum
Coordinate transformations: idx2crd, crd2idx
Layout operations: compact_order, prefix_product
Structural comparisons: congruent, compatible, weakly_congruent

Example usage:

from layout import IntTuple
from layout.int_tuple import flatten, compact_order, size

# Create nested tuples
var shape = IntTuple(2, IntTuple(3, 4), 5)  # Represents shape (2, (3, 4), 5)

# Flatten a nested tuple
var flat = flatten(shape)  # Results in (2, 3, 4, 5)

# Create compact strides for a given shape and order
var order = IntTuple(1, IntTuple(2, 3), 4)
var strides = compact_order(shape, order)  # Results in (1, (2, 6), 24)

# Calculate total size (product of all elements)
var total_size = size(shape)  # Results in 120
from layout import IntTuple
from layout.int_tuple import flatten, compact_order, size

# Create nested tuples
var shape = IntTuple(2, IntTuple(3, 4), 5)  # Represents shape (2, (3, 4), 5)

# Flatten a nested tuple
var flat = flatten(shape)  # Results in (2, 3, 4, 5)

# Create compact strides for a given shape and order
var order = IntTuple(1, IntTuple(2, 3), 4)
var strides = compact_order(shape, order)  # Results in (1, (2, 6), 24)

# Calculate total size (product of all elements)
var total_size = size(shape)  # Results in 120

Aliases

INT_TUPLE_VALIDATION = False:
IntList = List[Int, True]: A type alias for a List of integers with ownership. This alias defines a List that contains Int values and has ownership of its data. It's used throughout the module for storing and manipulating collections of integers, particularly for operations like permutations and indices.
UNKNOWN_VALUE = -1: Special value indicating an unknown or unspecified dimension. This constant is used throughout the IntTuple system to represent dimensions that are not known at compile time or have not been specified.

Structs

IntArray: A memory-efficient, register-passable array of integers.
IntTuple: A hierarchical, nested tuple of integers with efficient memory management.

Functions

abs: Compute the absolute value of each element in an IntTuple.
apply: Apply a function to each integer value in an IntTuple.
apply_predicate: Apply a predicate function recursively to two IntTuples.
apply_zip: Apply a function to pairs of elements from two IntTuples.
compact_order: Create a compact stride based on shape and order.
compatible: Test if two shapes are compatible for tensor operations.
congruent: Test if two IntTuples have the same hierarchical structure.
crd2idx: Map a logical coordinate to a linear index.
depth: Calculates the maximum nesting depth of an IntTuple.
fill_like: Creates an IntTuple with the same structure as the source but filled with a specified value.
flatten: Flatten a nested IntTuple into a single-level IntTuple.
idx2crd: Converts a linear index to a coordinate tuple within a given shape.
idx2crd2: Convert a linear index to coordinates.
inner_product: Compute the inner product of two IntTuples.
is_int: Check if an IntTuple represents a single integer value.
is_tuple: Check if an IntTuple represents a nested tuple.
lt: Compare two IntTuples lexicographically.
mul: Multiply each element in an IntTuple by a scalar value.
prefix_product: Compute the exclusive prefix product of an IntTuple.
product: Calculate the product of all values in an IntTuple.
product_each: Compute the product of elements in each sub-tuple of an IntTuple.
propagate_unknown: Propagates unknown dimensions from the target IntTuple to the source IntTuple.
reduce: Apply a reduction function to an IntTuple with an initial value.
reverse: Reverses the order of elements in an IntTuple, recursively.
shallow_apply: Apply a function to each top-level element of an IntTuple.
shape_div: Performs division operation between shape tuples.
signum: Calculate the sign of an integer.
size: Calculate the total size (product of all elements) of an IntTuple.
sorted: Sort an IntTuple using the provided comparison function.
sum: Calculate the sum of all values in an IntTuple.
to_nest: Nests a flat IntTuple according to the structure of a nested IntTuple.
to_unknown: Create an IntTuple with the same structure but filled with UNKNOWN_VALUE.
tuple_max: Calculate the maximum value in an IntTuple.
tuple_min: Compute the element-wise minimum of two IntTuples.
weakly_compatible: Test if shape A is weakly compatible with shape B.
weakly_congruent: Test if two IntTuples have similar hierarchical structures.
zip: Create a zip iterator from an array of IntTuple pointers.

Aliases​

Structs​

Functions​

Aliases

Structs

Functions