latest (dev)
Copy
Latest development documentation · Updated 2026-10-15
NestedIOFunction
- class tensorplay.autograd.function.NestedIOFunction[source]
Kept only for import compatibility; the modern contract is to define
forward+backwardonFunctiondirectly.- classmethod apply(*args, **kwargs)
Runs the operation and attaches gradient bookkeeping to outputs.
flat arguments computes
needs_input_gradand wires next-edges BEFORE forward; outputs are marked and attached AFTERsetup_context. When the fused C++ apply is present the hot path makes a single pybind crossing that also wires the backward entry; otherwise a generic Python fallback runs.
- static jvp(ctx, *grad_inputs)
Defines a formula for computing the jacobian-vector product.
Called by forward-mode AD when at least one input of this Function carries a tangent:
grad_inputsholds the tangent of each input (Nonefor inputs without one), and the returned tangents are attached to the outputs one-for-one.
- static setup_context(ctx, inputs, output)
Sets up the context object (Usage 2 above).
- Parameters:
ctx (_Context) – context object to modify in-place
inputs (tuple) – inputs to
forward()output (Any) – output of
forward()
- static vmap(info, in_dims, *args)
Defines a formula for vectorizing the operation.
Not yet supported by this engine; provided for API compatibility.
Help improve this page
Found an error, an unclear step, or a missing example?

