warp.Tape

class warp.Tape

Record kernel launches within a Tape scope to enable automatic differentiation. Gradients can be computed after the operations have been recorded on the tape via Tape.backward().

Example

tape = warp.Tape()

# forward pass
with tape:
    warp.launch(kernel=compute1, inputs=[a, b], device="cuda")
    warp.launch(kernel=compute2, inputs=[c, d], device="cuda")
    warp.launch(kernel=loss, inputs=[d, l], device="cuda")

# reverse pass
tape.backward(l)

Gradients can be accessed via the tape.gradients dictionary, e.g.:

print(tape.gradients[a])

Methods

backward([loss, grads])	Evaluate the backward pass of the recorded operations on the tape.
get_adjoint(a)	Return the adjoint container for a kernel argument.
record_func(backward, arrays)	Records a custom function to be executed only in the backward pass.
record_launch(kernel, dim, max_blocks, ...)
record_scope_begin(scope_name[, metadata])	Begin a scope on the tape to group operations together.
record_scope_end([remove_scope_if_empty])	End a scope on the tape.
reset()	Clear all operations recorded on the tape and zero out all gradients.
visualize([filename, simplify_graph, ...])	Visualize the recorded operations on the tape as a GraphViz diagram.
zero()	Zero out all gradients recorded on the tape.

backward(loss=None, grads=None)

Evaluate the backward pass of the recorded operations on the tape.

A single-element array loss or a dictionary of arrays grads can be provided to assign the incoming gradients for the reverse-mode automatic differentiation pass. If a grads entry targets an array with no .grad buffer, the incoming gradient array is assigned directly as that array’s .grad. If the target array already has a .grad buffer, the incoming gradient is copied into that existing buffer. When passing externally owned buffers, such as PyTorch grad_output tensors, allocate or attach an independent output gradient buffer first if Warp should not retain the external buffer.

Parameters:

• loss (array | None) – A single-element array that holds the loss function value whose gradient is to be computed

• grads (dict[array, array] | None) – A dictionary of arrays that map from Warp arrays to their incoming gradients

Note

When wp.config.verify_autograd_array_access is enabled, the read flags of the arrays recorded on this tape are cleared at the end of the backward pass, since subsequent writes can no longer corrupt this tape’s gradients. The flags are shared per array, so writes are then also no longer flagged against another tape that read the same arrays and has not yet run its own backward pass, or against a second backward() call on this tape.

record_launch(

kernel,

dim,

max_blocks,

inputs,

outputs,

device,

block_dim=0,

metadata=None,

)

record_func(backward, arrays)

Records a custom function to be executed only in the backward pass.

Parameters:

• backward (Callable) – A callable Python object (can be any function) that will be executed in the backward pass.

• arrays (list) – A list of arrays that are used by the backward function. The tape keeps track of these to be able to zero their gradients in Tape.zero()

record_scope_begin(scope_name, metadata=None)

Begin a scope on the tape to group operations together. Scopes are only used in the visualization functions.

record_scope_end(remove_scope_if_empty=True)

End a scope on the tape.

Parameters:

remove_scope_if_empty (bool) – If True, the scope will be removed if no kernel launches were recorded within it.

get_adjoint(a)

Return the adjoint container for a kernel argument.

For warp.array inputs with gradients enabled, this returns warp.array.grad and tracks it in gradients so it can be zeroed later. For instances created with warp.struct(), a mirrored struct is created with adjoints for array fields, and nested structs are handled recursively. Registered native values receive a default-initialized value so reverse-launch argument packing preserves their ABI.

Parameters:

a – Kernel argument to map to an adjoint. Can be a warp.array, a warp.struct() instance, or a value type.

Returns:

The adjoint object for a or None if no adjoint is required.

reset()

Clear all operations recorded on the tape and zero out all gradients.

zero()

Zero out all gradients recorded on the tape.

visualize(

filename=None,

simplify_graph=True,

hide_readonly_arrays=False,

array_labels=None,

choose_longest_node_name=True,

ignore_graph_scopes=False,

track_inputs=None,

track_outputs=None,

track_input_names=None,

track_output_names=None,

graph_direction='LR',

)

Visualize the recorded operations on the tape as a GraphViz diagram.

Example

import warp as wp

tape = warp.Tape()
with tape:
    # record Warp kernel launches here
    warp.launch(...)

dot_code = tape.visualize("tape.dot")

This function creates a GraphViz dot file that can be rendered into an image using the GraphViz command line tool, e.g. via

dot -Tpng tape.dot -o tape.png

Parameters:

• filename (str | None) – The filename to save the visualization to (optional).

• simplify_graph (bool) – If True, simplify the graph by detecting repeated kernel launch sequences and summarizing them in subgraphs.

• hide_readonly_arrays (bool) – If True, hide arrays that are not modified by any kernel launch.

• array_labels (dict[array, str] | None) – A dictionary mapping arrays to custom labels.

• choose_longest_node_name (bool) – If True, the automatic name resolution will aim to find the longest name for each array in the computation graph.

• ignore_graph_scopes (bool) – If True, ignore the scopes recorded on the tape when visualizing the graph.

• track_inputs (list[array] | None) – A list of arrays to track as inputs in the graph to ensure they are shown regardless of the hide_readonly_arrays setting.

• track_outputs (list[array] | None) – A list of arrays to track as outputs in the graph so that they remain visible.

• track_input_names (list[str] | None) – A list of custom names for the input arrays to track in the graph (used in conjunction with track_inputs).

• track_output_names (list[str] | None) – A list of custom names for the output arrays to track in the graph (used in conjunction with track_outputs).

• graph_direction (str) – The direction of the graph layout (default: “LR”).

Returns:

The dot code representing the graph.

Return type:

str