a + b, End to End

a + b, End to End A sequence diagram generated by Archify. c = a + b _C_engine.add _unwrap_or_scalar turns b into a TensorImpl and _maybe_promote aligns dtypes before the call forward(a, b) add_op goes straight into BinaryKernel<AddBackward>::forward SchemaGuard Dtype and device mismatches are rejected first; then determinism is checked and AMP picks the autocast dtype eff_dt for_device(dev) Indexes an array by the Device enum; 'metal' maps to Device.GPU CPU: vDSP_vadd Takes the output buffer from the pool and computes synchronously GPU: mlx add mlx::core::add only adds a graph node; the math runs at eval() or item() Storage new TensorImpl Wraps the Storage with the broadcast shape; no AutogradMeta yet AddBackward node Only when GradMode is on and an input requires grad set_grad_fn Also records next_edges and saved_versions TensorImplPtr Crosses pybind11 as a shared_ptr, no copy _wrap → Tensor Python Validate Compute Record Return User · c = a + b · Sequence participant User c = a + b __add__ · Tensor · _dunders · Sequence participant __add__ Tensor · _dunders pybind11 · _C_engine.add · Sequence participant pybind11 _C_engine.add kernel · AddBackward · Sequence participant kernel AddBackward core · SchemaGuard · Sequence participant core SchemaGuard backend · Cpu/GpuBackend · Sequence participant backend Cpu/GpuBackend Apple · vDSP · MLX · Sequence participant Apple vDSP · MLX autograd · AutogradMeta · Sequence participant autograd AutogradMeta Legend request return guard check lazy or conditional default message

One object, three roles

  • • AddBackward is the op, the BinaryKernel (CRTP) and the autograd node at once
  • • a + b skips the _ops registry: the injected __add__ calls the engine directly

CPU vs GPU

  • • CPU: vDSP_vadd computes on the spot
  • • GPU: mlx::core::add only adds a node; it runs at .item() or eval()

Where it is rejected

  • • DtypeMismatch, DeviceMismatch, ShapeMismatch
  • • C++ never promotes dtypes; Python aligns them first