package nx

  1. Overview
  2. Docs
package nx
Legend:
Page
Library
Module
Module type
Parameter
Class
Class type
Source

Module Nx_core.Symbolic_shapeSource

Symbolic shapes for shape-polymorphic tensors.

This module provides symbolic dimensions that can be bound at runtime, enabling dynamic shapes in compiled kernels and graph optimization.

Sourcetype var

A symbolic variable representing a dimension

Sourcetype expr =
  1. | Const of int
  2. | Var of var
  3. | Add of expr * expr
  4. | Mul of expr * expr
  5. | Neg of expr

Expression type for dimension arithmetic

Sourcetype dim = expr

A dimension is an expression

Sourcetype t = dim array

A shape is an array of dimensions

Creation

Sourceval static : int -> dim

static n creates a static dimension with value n.

Sourceval dynamic : string -> min:int -> max:int -> dim

dynamic name ~min ~max creates a dynamic dimension with bounds.

Sourceval of_ints : int array -> t

of_ints arr creates a shape with all static dimensions.

Sourceval of_list : int list -> t

of_list lst creates a shape with all static dimensions.

Dimension Operations

Sourceval add : dim -> dim -> dim

add d1 d2 creates a dimension expression d1 + d2.

Sourceval mul : dim -> dim -> dim

mul d1 d2 creates a dimension expression d1 * d2.

Sourceval neg : dim -> dim

neg d creates a dimension expression -d.

Runtime Binding

Sourceval bind : string -> int -> t -> unit

bind name value shape binds a concrete value to all variables with name in shape.

Sourceval eval : t -> int array option

eval shape returns concrete shape if all dimensions are bound.

Sourceval eval_dim : dim -> int option

eval_dim dim evaluates a single dimension.

Sourceval partial_eval : t -> int option array

partial_eval shape returns an array of evaluated dimensions, with None for unbound.

Sourceval is_fully_bound : t -> bool

is_fully_bound shape returns true if all symbolic dimensions are bound.

Sourceval numel : t -> int option

numel shape returns the total number of elements if shape is fully bound.

Special Values

Sourceval infer : dim

Special dimension value representing "infer from context" (like -1 in NumPy reshape).

Sourceval is_infer : dim -> bool

is_infer dim returns true if dimension should be inferred.

Shape Resolution

Sourceval resolve_reshape : from_shape:t -> to_shape:t -> t option

resolve_reshape ~from_shape ~to_shape resolves a reshape operation, computing any infer dimensions in to_shape based on from_shape's numel. Returns None if resolution is impossible.

Sourceval substitute : (string * int) list -> t -> t

substitute bindings shape substitutes variable bindings into a shape, replacing variables with their bound values.

Analysis

Sourceval vars : t -> var list

vars shape returns all unique symbolic variables in shape.

Sourceval is_static : t -> bool

is_static shape returns true if all dimensions are static.

Sourceval rank : t -> int

rank shape returns number of dimensions.

Utilities

Sourceval to_string : t -> string

to_string shape returns human-readable representation.

Sourceval equal : t -> t -> bool

equal s1 s2 compares shapes structurally.