CRTP base for single-input, single-output op kernels.
Inherits AutogradNode\<Derived, 1\> (one saved input slot)
and kernel::IKernel so an instance can be held polymorphically
while still exposing the typed static forward() trampoline.
Concrete ops declare themselves as
class FooBackward : public UnaryKernel<FooBackward> and provide:
static constexpr OpSchema schema_v1— op name + AMP policy.static cpu_kernel(a, out_shape, dtype) -> CpuStorageand/orstatic gpu_kernel(a, out_shape, dtype) -> GpuStorage, orstatic dispatch(IBackend&, a, out_shape, dtype) -> Storage.Storage grad_formula(Storage grad_out)— the local Jacobian product computed at backward time.- Optional
grad_formula_impl(grad_out, a_impl, out_impl) -> TensorImplPtrfor graph-mode (create_graph=True) higher-order differentiation.
Template Parameters
Derived : class The concrete CRTP self-type.
See Also
BinaryKernel, NaryKernel, VariadicKernel.
IKernel — the abstract base above the CRTP layer.
Attributes
kSavesInputstatic constexpr boolforward() snapshots a->storage() into saved_inputs_[0] for use in grad_formula. Defaults to true; set false when the gradient formula uses only the saved output (e.g. ReLU keyed on output sign).kSavesOutputstatic constexpr boolforward() additionally snapshots the output storage into saved_output_. Defaults to false. Set true for ops whose backward is cheap in terms of y (e.g. Exp, Sigmoid).kHasGradientstatic constexpr booltrue; set false for non-differentiable ops to skip all graph wiring.Notes
AMP policy and dtype promotion are forwarded through
SchemaGuard using Derived::schema_v1; the effective dtype
returned drives both the input cast and the saved-input dtype.
Static methods
1int forward(const int & a)Typed forward trampoline for a single-input op.
Parameters
aconst std::shared_ptr<TensorImpl>&Returns
std::shared_ptr<TensorImpl>The output tensor. If a requires a gradient (and grad mode is enabled) the result has its grad_fn set to a freshly constructed Derived backward node.
Raises
ShapeMismatchDtypeMismatchNotes
Behaviour
- Validate
ais non-null. - Resolve the effective dtype via
SchemaGuard. - Materialise contiguous on CPU when
ais non-contiguous. - Cast to the effective dtype if needed.
- Dispatch to
dispatch/gpu_kernel/cpu_kernel(see file header for priority). - Wire the autograd graph when
kHasGradientand any input requires a gradient.
Methods
5Backward implementation invoked by the autograd engine.
Parameters
grad_outStorageReturns
std::vector<Storage>Single-element vector containing grad_input reduced back to input_shapes_[0] via reduce_grad_to_shape.
Notes
The reduction step is needed when the forward broadcast the input to a larger output shape; for purely shape-preserving unary ops it is a no-op.
Graph-mode backward — supports create_graph=True.
Parameters
grad_outconst TensorImplPtr&TensorImpl with its own grad_fn for higher-order differentiation.Returns
std::vector<TensorImplPtr>Single-element vector containing the fully-traced grad_input reduced back to the original input shape via sum_op / reshape_op.
Raises
:runtime_errorgrad_formula_impl.grad_formula_impl
→TensorImplPtrint grad_formula_impl(const int &, const int &, const int &)Default graph-mode gradient formula — concrete ops override.
Parameters
gconst TensorImplPtr&TensorImpl that itself carries grad_fn so the backward computation is differentiable.aconst TensorImplPtr&outconst TensorImplPtr&kSavesOutput == true).Returns
TensorImplPtrgrad_input as a fully-traced TensorImpl.
Raises
:runtime_errorcreate_graph=True.Return the canonical schema name of the concrete op.
Return the autograd node label (same string as name).