tensorplay.autograd.function API

Records operation history and defines formulas for differentiating ops.

1. Legacy style: forward(ctx, ...) / backward(ctx, ...) (forward receives a context object).

2. Combined-forward style: define forward(*args, **kwargs), setup_context(ctx, inputs, output) and use save_for_backward/save_for_forward inside setup_context instead of receiving a ctx argument in forward.

classmethod apply(*args, **kwargs)[source]

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 backward(ctx, *grad_outputs)[source]

Defines a formula for differentiating the operation.

static forward(ctx, *args, **kwargs)[source]

Performs the operation.

This function is to be overridden by all subclasses. There are two ways to define forward:

Usage 1 (Combined forward and ctx):

@staticmethod
def forward(ctx, input1, input2):
    ...
    return output

Usage 2 (Separated forward and ctx):

@staticmethod
def forward(input1, input2):
    ...
    return output

@staticmethod
def setup_context(ctx, inputs, output):
    ...

static jvp(ctx, *grad_inputs)[source]

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)[source]

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)[source]

Defines a formula for vectorizing the operation.

Not yet supported by this engine; provided for API compatibility.

the name classproperty ("<Cls>Backward", used for node naming) and a friendlier repr for subclasses.

mro()

Return a type’s method resolution order.

In-place operations must call ctx.mark_dirty on the mutated inputs inside forward; this subclass exists only so historical code that subclasses it keeps working.

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 backward(ctx, *grad_outputs)

Defines a formula for differentiating the operation.

static forward(ctx, *args, **kwargs)

Performs the operation.

This function is to be overridden by all subclasses. There are two ways to define forward:

Usage 1 (Combined forward and ctx):

@staticmethod
def forward(ctx, input1, input2):
    ...
    return output

Usage 2 (Separated forward and ctx):

@staticmethod
def forward(input1, input2):
    ...
    return output

@staticmethod
def setup_context(ctx, inputs, output):
    ...

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.

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.