TensorPlay
latest (dev)
Copy
View Markdown

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 + backward on Function directly.

classmethod apply(*args, **kwargs)

Runs the operation and attaches gradient bookkeeping to outputs.

flat arguments computes needs_input_grad and wires next-edges BEFORE forward; outputs are marked and attached AFTER setup_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_inputs holds the tangent of each input (None for 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.

On this page

Ask DeepWiki