warp.Tape#

class warp.Tape[source]#

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 = wp.Tape()

# forward pass
with tape:
    wp.launch(kernel=compute1, inputs=[a, b], device="cuda")
    wp.launch(kernel=compute2, inputs=[c, d], device="cuda")
    wp.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])

__init__()[source]#

Methods

`__init__`()
`backward`([loss, grads])	Evaluate the backward pass of the recorded operations on the tape.
`get_adjoint`(a)
`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)[source]#

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.

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

record_launch( kernel, dim, max_blocks, inputs, outputs, device, block_dim=0, metadata=None, )[source]#

record_func(backward, arrays)[source]#

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

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

reset()[source]#: Clear all operations recorded on the tape and zero out all gradients.

zero()[source]#: 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', )[source]#

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

Example

import warp as wp

tape = wp.Tape()
with tape:
    # record Warp kernel launches here
    wp.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