Purpose: This module defines all mathematical operations and the
Operation.__call__()pipelineβthe heart of Nabla's execution model.
When you call any operation (e.g., add(x, y) or x + y), Operation.__call__() in base.py orchestrates a 9-step pipeline. The key architectural insight: steps 1-3 and 6-9 always run (metadata + tracing), while steps 4 and 5 are conditional.
βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
β Operation.__call__() - The 9-Step Pipeline β
β β
β βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β STEP 1: METADATA COLLECTION (base.py ~line 101) β β
β β β β
β β max_batch_dims, any_traced, any_sharded, any_has_tangent = β β
β β collect_metadata(args) β β
β β β β
β β Purpose: Scan inputs to determine execution context β β
β β β’ max_batch_dims: for vmap axis translation β β
β β β’ any_traced: whether to record OpNode for autodiff β β
β β β’ any_sharded: whether SPMD logic needed β β
β β β’ any_has_tangent: whether JVP propagation needed β β
β βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β β
β βΌ β
β βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β STEP 2: ADAPTATION & RESHARDING (base.py ~line 104-105) β β
β β β β
β β resharded_args, adapted_kwargs, predicted_output_spec, mesh, β β
β β reduce_axes = adapt_and_reshard(self, args, kwargs, ...) β β
β β β β
β β Internally: β β
β β β’ adapt_kwargs(): translate logicalβphysical (axis by batch_dims) β β
β β β’ infer_output_sharding(): predict output sharding + reductions β β
β β β’ reshard_inputs(): insert AllGather/AllToAll if needed β β
β β β β
β β Key outputs: β β
β β β’ reduce_axes: mesh axes needing post-op reduction (step 8) β β
β β β’ resharded_args: inputs with correct shardings for execution β β
β βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β β
β βΌ β
β βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β STEP 3: COMPUTE STRUCTURAL HASH (base.py ~line 108) β β
β β β β
β β op_hash = compute_structural_hash(self.name, resharded_args, β β
β β adapted_kwargs) β β
β β β β
β β Purpose: Create cache key for compiled model lookup β β
β β β’ Unrealized tensors: keyed by their own (op_hash, output_index) β β
β β β’ Realized tensors: keyed by (dtype, shape, sharding) β β
β β This enables GRAPH.evaluate() to skip graph building on cache hit β β
β βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β β
β βΌ β
β βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β STEP 4: EAGER EXECUTION (CONDITIONAL) β‘ (base.py ~line 116) β β
β β β β
β β execution_results = eager_execute(self, resharded_args, kwargs, β β
β β adapted_kwargs) β β
β β β β
β β if EAGER_MAX_GRAPH=1: β β
β β β Calls op.execute(resharded_args, kwargs) β β
β β β Returns (shard_graph_values, output_sharding, mesh) β β
β β β β
β β if EAGER_MAX_GRAPH=0 (default): β β
β β β Returns None (graph building DEFERRED) β β
β βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β β
β βΌ β
β βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β STEP 5: ANALYTICAL METADATA (CONDITIONAL) (base.py ~line 118) β β
β β β β
β β if execution_results is None or VERIFY_EAGER_SHAPES: β β
β β output_physical_shapes, output_shard_dtypes, ... = β β
β β self.compute_physical_shape(...) β β
β β else: β β
β β β Set to None (Metadata trusted from backend values) β β
β β β β
β β Purpose: Manually determine shapes if graph build is deferred β β
β β or cross-verify backend results if VERIFY_EAGER_SHAPES=1. β β
β βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β β
β βΌ β
β βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β STEP 6: PACKAGING (base.py ~line 148) β β
β β β β
β β output = package_outputs(self, execution_results, β β
β β output_physical_shapes, output_shard_dtypes, ...) β β
β β β β
β β Creates Tensor with metadata: β β
β β β’ if step 5 set metadata β Metadata stored in TensorImpl β β
β β β’ if step 5 skipped β Metadata lazily retrieved from backend valuesβ β
β β β β
β β if EAGER_MAX_GRAPH (step 4 ran): β β
β β output._impl._graph_values = [TensorValue, ...] β β
β β β β
β β if DEFERRED (step 4 returned None): β β
β β output._impl.graph_values_epoch = -1 β "PROMISE TENSOR" β β
β βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β β
β βΌ β
β βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β STEP 7: SETUP OUTPUT REFS (OpNode Creation) (base.py ~line 129) β β
β β β β
β β self._setup_output_refs(output, resharded_args, kwargs, β β
β β op_hash=op_hash) β β
β β β β
β β Creates OpNode with: β β
β β β’ _refs: tuple of output TensorImpls β β
β β β’ op: the operation instance β β
β β β’ op_args: input TensorImpls (for backward traversal) β β
β β β’ op_kwargs: ORIGINAL kwargs (critical for rehydration!) β β
β β β’ _op_hash: structural hash for caching β β
β β β β
β β Key: Stores ORIGINAL kwargs, not adapted. execute() adapts β β
β β internally, so rehydration can replay correctly. β β
β βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β β
β βΌ β
β βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β STEP 8: AUTO-REDUCTION (base.py ~line 130) β β
β β β β
β β output = apply_auto_reduction(self, output, mesh, reduce_axes) β β
β β β β
β β If contracting factors (like K in matmul) were sharded: β β
β β β Partial sums exist across devices β β
β β β Inserts AllReduce with op.collective_reduce_type (sum/max/...) β β
β β β Updates sharding spec to remove reduced axes β β
β βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β β
β βΌ β
β βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β STEP 9: JVP PROPAGATION (base.py ~line 132-133) β β
β β β β
β β if any_has_tangent: β β
β β apply_jvp(self, args, output) β β
β β β β
β β Forward-mode autodiff: propagate tangents via op.jvp_rule() β β
β β Attaches tangent to output._impl.tangent β β
β βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ β
β β β
β βΌ β
β return output β
βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
MAX graph construction has overhead. The default deferred mode optimizes for cache hits:
βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
β Why Defer Graph Building? β
βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ€
β β
β SCENARIO: Training loop iteration β
β βββββββββββββββββββββββββββββββββββ β
β for batch in dataloader: # Thousands of iterations β
β loss = model(batch) # Same computation structure β
β grads = grad(loss_fn)(params) # Same structure, different data β
β β
β WITH EAGER GRAPH BUILDING (EAGER_MAX_GRAPH=1): β
β β’ Every op.execute() builds MAX graph nodes β
β β’ evaluate() compiles the graph β
β β’ Cache stores compiled model β
β β’ Next iteration: graph built AGAIN, then cache hit skips compile β
β β’ Wasted work: graph building every iteration β
β β
β WITH DEFERRED GRAPH BUILDING (default): β
β β’ Operations only compute metadata (shapes, dtypes) β
β β’ evaluate() checks cache by op_hash β
β β’ Cache HIT β Skip graph building entirely, just run cached model β
β β’ Cache MISS β Build graph via _replay_trace_to_build_graph() β
β β’ Hot paths skip ALL graph construction overhead β
β β
βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
- User code needs
.shapefor control flow and debugging. - Sharding propagation needs shapes to plan data movement.
- Type checking and broadcasting validation must happen immediately.
In Eager Mode, we can skip this analytical step because the backend provides the shapes directly. This ensures that the Python-side view and backend-side reality never drift.
This is a critical design decision for trace rehydration:
# In __call__:
raw_result = self.execute(resharded_args, kwargs) # β original kwargs!
# Inside execute (default implementation):
def execute(self, args, kwargs):
adapted_kwargs = self.adapt_kwargs(args, kwargs, batch_dims) # adapts here
shard_results = spmd.execute_on_shards(self.kernel, args, adapted_kwargs, mesh)During rehydration (replaying the trace for backward pass), we only have access to the original kwargs stored in OpNode. If __call__ passed pre-adapted kwargs to execute, rehydration would fail or produce wrong results.
Operation (base class - defines __call__ pipeline)
β
βββ UnaryOperation # Single input: relu, exp, neg
β βββ Default execute loops over shards
β
βββ BinaryOperation # Two inputs: add, mul, matmul
β βββ Overrides __call__ to broadcast shapes BEFORE parent pipeline
β
βββ LogicalAxisOperation # Ops with axis kwargs: reduce, transpose
β βββ Implements adapt_kwargs() to translate axis by batch_dims
β βββ ReduceOperation # reduce_sum, mean
β βββ LogicalShapeOperation # reshape, transpose
β
βββ CollectiveOperation # Communication ops with custom execution
| Method | Required | Purpose | Notes |
|---|---|---|---|
name |
β | String identifier | Property |
kernel(*args, **kwargs) |
β | MAX primitive (per-shard) | Called by execute() in shard loop |
compute_physical_shape(args, kwargs, output_sharding) |
β NEW | Infer output shapes WITHOUT building MAX graph | Must return (shapes, dtypes, devices) |
execute(args, kwargs) |
Has default | Custom execution logic | Override if needed. Receives ORIGINAL kwargs! |
adapt_kwargs(args, kwargs, batch_dims) |
Has default | Translate logicalβphysical kwargs | Override for axis ops |
vjp_rule(primals, cotangent, output) |
For autodiff | Backward gradient | See core/autograd/ |
jvp_rule(primals, tangents, output) |
For forward-mode | Forward tangent | utils.py |
sharding_rule(input_shapes, output_shapes) |
For SPMD | Factor-based sharding | propagation.py |
collective_reduce_type |
Has default="sum" | Reduction op type | For apply_auto_reduction |
Every operation must implement compute_physical_shape. This is how shapes are computed eagerly even when graph building is deferred:
def compute_physical_shape(
self, args: tuple, kwargs: dict, output_sharding: Any = None
) -> tuple[list[tuple[int, ...]], list[DType], list[Device]]:
"""
Infer per-shard physical shapes, dtypes, and devices for outputs.
CRITICAL: Must NOT build MAX graph nodes! Only compute metadata.
"""Optimization Path: If EAGER_MAX_GRAPH=1, this method is only called if VERIFY_EAGER_SHAPES=1 is enabled for cross-verification. Otherwise, Nabla trusts the shapes returned by the MAX backend directly, which is faster and more reliable. In Deferred mode (default), this method remains the primary way to compute shapes.
y = matmul(A, B) # A: [M, K] sharded on K, B: [K, N]Step 1 (Metadata): Collects batch_dims=0, traced=True, sharded=True
Step 2 (Adaptation):
infer_output_shardinguses rule"m k, k n -> m n"- Factor
kis sharded β appears in reduce_axes - If B's K dim has different sharding β
reshard_inputsinserts communication
Step 3 (Hash): op_hash = ("matmul", (A_key, B_key), kwargs_key)
Step 4 (Eager Execution): EAGER_MAX_GRAPH=0 β Returns None (no graph building!)
Step 5 (Physical Shapes): compute_physical_shape returns ([M, N], [M, N], ...) per shard
Step 6 (Packaging):
- Creates output tensor with
_physical_shapes=[(M,N), ...] - Sets
graph_values_epoch = -1(promise tensor) - Calls
GRAPH.add_unrealized(output._impl)
Step 7 (OpNode): Records OpNode with op_hash for cache lookup
Step 8 (Auto-Reduction): reduce_axes={'k_axis'} β Inserts AllReduce operation
Step 9 (JVP): Not applicable here
Later, when y.numpy() is called:
GRAPH.evaluate(y)checks cache byop_hash- Cache MISS:
_replay_trace_to_build_graph()callsmatmul.execute() - Compiles and runs, stores to
y._impl._buffers - Caches compiled model for next time
| File | Purpose | Key Exports |
|---|---|---|
| base.py | __call__ pipeline, base classes |
Operation, BinaryOperation, UnaryOperation, ReduceOperation, LogicalAxisOperation |
| utils.py | Execution helpers | eager_execute, package_outputs, collect_metadata, apply_auto_reduction, apply_jvp |
| binary.py | Binary ops | add, sub, mul, div, matmul, pow |
| unary.py | Unary ops | relu, sigmoid, tanh, exp, log, neg, softmax |
| reduction.py | Reductions | reduce_sum, mean, reduce_max, reduce_min |
| creation.py | Tensor factories | full, zeros, ones, arange, uniform, gaussian |
| view/ | Shape ops | reshape, transpose, squeeze, broadcast_to, gather, scatter |
| communication/ | Collectives | all_reduce, all_gather, shard, reshard, reduce_scatter |
| comparison.py | Comparisons | equal, not_equal, greater, less |
| control_flow.py | Control flow | where, cond, while_loop, scan |
| multi_output.py | Multi-output | split, chunk, unbind |
| custom_op.py | Extensions | call_custom_kernel |
AI Agents - Critical Rules:
compute_physical_shaperequired: Every op must implement this. It cannot build graph nodes.executecontract: MUST receive original kwargs and adapt internally. Breaking this breaks rehydration._setup_output_refs: Storesresharded_args, not original args. Backward pass needs correctly sharded inputs.- OpNode stores original kwargs: For rehydration correctness.
- Testing: Every op needs gradient tests via
nabla.grad(fn)(x).