CRTP base for two-input, single-output op kernels — pairs forward output computation with the two-input saved-tensor bookkeeping every binary op needs.
Inherits AutogradNode\<Derived, 2\> (two saved input slots)
and kernel::IKernel. forward owns all validation,
broadcast inference, dispatch, and graph wiring; apply calls
Derived::grad_formula and broadcast-reduces each gradient back to
its original input shape.
Concrete ops declare themselves as
class FooBackward : public BinaryKernel<FooBackward> and provide:
static constexpr OpSchema schema_v1— op name + AMP policy.static cpu_kernel(a, b, out_shape, dtype) -> CpuStorageand/orstatic gpu_kernel(a, b, out_shape, dtype) -> GpuStorage, orstatic dispatch(IBackend&, a, b, out_shape, dtype) -> Storage.grad_formula(grad_out) -> std::tuple<Storage, Storage>for the local Jacobian product.- Optional
grad_formula_impl(grad_out, a_impl, b_impl) -> std::pair<TensorImplPtr, TensorImplPtr>for create_graph higher-order differentiation.
Template Parameters
Derived : class The concrete CRTP self-type.
See Also
UnaryKernel, NaryKernel, VariadicKernel.
IKernel — the abstract base above the CRTP layer.
Attributes
kSavesInputsstatic constexpr boolforward() snapshots both input storages into saved_inputs_. Defaults to true; set false for ops whose backward does not consult the original inputs (e.g. Add whose backward is identity).Notes
Slot count is 2 (one per input). Mixed-dtype binary ops are not
implicitly promoted — callers must cast. Cross-device operands are
never silently moved — callers must migrate. Both checks throw
DtypeMismatch / DeviceMismatch from forward.
Static methods
3Convenience accessor for the backend matching a device.
int forward(const int & a, const int & b)Typed forward trampoline for a two-input op.
Parameters
a, bconst std::shared_ptr<TensorImpl>&Returns
std::shared_ptr<TensorImpl>The output tensor. If grad mode is on and either input requires a gradient the result carries a fresh Derived grad_fn with two outbound edges.
Raises
DtypeMismatcha->dtype() != b->dtype().DeviceMismatcha->device() != b->device().ShapeMismatchreduce_impl_to_shape
→TensorImplPtrint reduce_impl_to_shape(const int & grad, const int & grad_shape, const int & target_shape)Reduce a traced gradient tensor back to a target input shape.
Parameters
gradconst TensorImplPtr&grad_shape.grad_shapeconst Shape&grad.target_shapeconst Shape&Returns
TensorImplPtrA new tensor whose shape equals target_shape. Sums over leading broadcast axes and over any axis where target_shape is unit but grad_shape is not, then reshapes to target_shape exactly.
Notes
Used inside apply_for_graph implementations and mirrors
the Storage-level helper reduce_grad_to_shape.
Methods
6Backward implementation invoked by the autograd engine.
Parameters
grad_outStorageReturns
std::vector<Storage>{grad_a, grad_b} each reduced back to its original input shape via reduce_grad_to_shape.
Notes
Reduction is needed because forward() may have broadcast
either operand to a larger output shape.
Graph-mode backward — supports create_graph=True.
Parameters
grad_outconst TensorImplPtr&TensorImpl with grad_fn.Returns
std::vector<TensorImplPtr>{grad_a, grad_b} as fully-traced TensorImpl pointers, each reduced to its original input shape via sum_op / reshape_op.
Raises
:runtime_errorDerived did not override grad_formula_impl.grad_formula_impl
→std::pair<TensorImplPtr, TensorImplPtr>int grad_formula_impl(const int &, const int &, const int &)Default graph-mode gradient formula — concrete ops override.
Parameters
grad_outconst TensorImplPtr&a, bconst TensorImplPtr&out_shape_.Returns
std::pair<TensorImplPtr, TensorImplPtr>(grad_a, grad_b) as fully-traced tensors.
Raises
:runtime_errorcreate_graph=True.Return the canonical schema name of the concrete op.
Return the autograd node label (same string as name).
Return the saved input k broadcast to out_shape_.
Parameters
kstd::size_tsaved_inputs_ (0 or 1).Returns
StorageThe saved storage when input_shapes_[k] == out_shape_, otherwise a freshly materialised broadcast copy. On GPU this calls mlx::core::broadcast_to; on CPU it calls detail::broadcast_cpu.
Notes
Lazy: the broadcast copy is only materialised when needed.
Used by grad_formula implementations that need the
full-rank operand rather than the possibly-smaller original.