Module Nx_core.ShapeSource

Shape operations for multi-dimensional arrays.

This module provides fundamental operations on array shapes including index conversions, broadcasting, and stride calculations.

Sourcetype t = int array

Shape representation as array of dimension sizes.

Basic Operations

Sourceval numel : t -> int

numel shape computes total number of elements.

Sourceval equal : t -> t -> bool

equal s1 s2 tests shape equality.

Strides

Sourceval c_contiguous_strides : t -> int array

c_contiguous_strides shape computes row-major strides.

Handles zero-size dimensions correctly.

Index Conversions

Sourceval ravel_index : int array -> int array -> int

ravel_index indices strides computes linear offset.

Sourceval unravel_index : int -> t -> int array

unravel_index offset shape converts to multi-dimensional indices.

Assumes C-contiguous layout.

Sourceval unravel_index_into : int -> t -> int array -> unit

unravel_index_into offset shape result writes indices into result array.

In-place version to avoid allocations in loops.

Shape Manipulation

Sourceval resolve_neg_one : t -> int array -> t

resolve_neg_one current_shape new_spec infers dimension marked with -1.

Sourceval broadcast : t -> t -> t

broadcast shape_a shape_b computes broadcast result shape.

Follows NumPy rules: dimensions match if equal or one is 1.

Sourceval broadcast_index : int array -> t -> int array

broadcast_index target_indices source_shape maps indices for broadcasting.

Returns indices in source shape corresponding to target position.

Sourceval broadcast_index_into : int array -> t -> int array -> unit

broadcast_index_into target_indices source_shape result writes broadcast indices into result array.

In-place version to avoid allocations in loops.

Pretty Printing

Sourceval pp : Format.formatter -> t -> unit

pp fmt shape prints shape in 2x3x4 format.

Sourceval to_string : t -> string

to_string shape converts to string "2x3x4".