nvalchemiops.torch.autograd: Autograd Utilities

Autograd Utilities for Warp-PyTorch Integration

This module provides utilities for integrating Warp’s automatic differentiation with PyTorch custom operators. It abstracts common patterns for:

1. Checking if any tensor requires gradients

2. Conditionally creating Warp tapes

3. Storing tape and warp arrays on output tensors

4. Retrieving them in backward passes

5. Decorator-based custom op registration with automatic backward generation

import warp as wp import torch from contextlib import contextmanager, nullcontext from typing import Any, Optional, Sequence, Union

Custom Op Registration

nvalchemiops.torch.autograd.warp_custom_op(name, outputs, grad_arrays=None, mutates_args=())

Decorator to create a Warp-backed PyTorch op with compile-safe autograd.

This decorator eliminates boilerplate by automatically generating: - A torch.library.custom_op forward registered with fake/meta support - A hidden token input for runtime state handoff while the public wrapper still exposes only the user-visible signature - A traceable register_autograd wrapper that replays Warp tapes through an opaque backward custom op - Stream binding so Warp launches execute on PyTorch’s current CUDA stream

Parameters:

• name (str) – Full custom op name (e.g., “alchemiops::_my_kernel”).

• outputs (list[OutputSpec]) – Specifications for each output tensor.

• grad_arrays (list[str], optional) – Names of warp arrays to track for gradients. Should include output names first, then differentiable input names. If None, auto-generated from outputs + all inputs that are likely differentiable (excludes common non-differentiable names like neighbor_list, batch_idx, etc.).

• mutates_args (tuple, default=()) – Arguments that are mutated by the op (passed to custom_op).

Returns:

Decorator function.

Return type:

Callable

Examples

>>> @warp_custom_op(
...     name="alchemiops::_ewald_real_space_energy",
...     outputs=[
...         OutputSpec("energies", wp.float64, lambda pos, *_: (pos.shape[0],)),
...     ],
...     grad_arrays=["energies", "positions", "charges", "cell", "alpha"],
... )
... def _ewald_real_space_energy(
...     positions: torch.Tensor,
...     charges: torch.Tensor,
...     cell: torch.Tensor,
...     alpha: torch.Tensor,
...     neighbor_list: torch.Tensor,
...     neighbor_shifts: torch.Tensor,
... ) -> torch.Tensor:
...     # Implementation here - no boilerplate needed!
...     ...
...     return energies

Notes

The decorated function should still call attach_for_backward() at the end of grad-enabled forward execution so the registered forward op can collect the runtime Warp tape and arrays from the real output tensor.

retain_graph=True is supported: the Warp tape is preserved across backward passes and zeroed before each replay. create_graph=True is not supported – Warp backward ops do not register a second-order autograd formula, so higher-order differentiation through them will raise. Use hybrid_forces=True in electrostatics APIs when you need to combine explicit Warp forces with autograd-based charge-gradient forces.

class nvalchemiops.torch.autograd.OutputSpec(name, dtype, shape, torch_dtype=None)

Specification for a custom op output.

Parameters:

• name (str) – Name of the output (used for backward pass).

• dtype (wp dtype) – Warp dtype (e.g., wp.float64, wp.vec3d).

• shape (Callable or tuple) – Either a tuple of ints, or a callable that takes the input tensors and returns the shape. For callable, signature should match the custom op’s input signature.

• torch_dtype (torch.dtype, optional) – PyTorch dtype override. If omitted (None), the dtype is inferred from the resolved Warp dtype via _wp_dtype_to_torch.

Examples

>>> OutputSpec("energies", wp.float64, lambda pos, *_: (pos.shape[0],))
>>> OutputSpec("forces", wp.vec3d, lambda pos, *_: (pos.shape[0], 3))
>>> OutputSpec("virial", wp.mat33d, (3, 3))  # Static shape

Warp-PyTorch Interop

nvalchemiops.torch.autograd.warp_stream_from_torch(*values)

Bind Warp launches to PyTorch’s current CUDA stream when tensors are CUDA.

Finds the first CUDA tensor in values and switches Warp to its stream for the duration of the context. Falls back to a no-op context when no CUDA tensor is found (e.g., CPU-only runs).

Parameters:

*values (Any) – Arbitrary arguments; only torch.Tensor instances on CUDA are inspected for stream information.

Yields:

torch.cuda.Stream or None – The active PyTorch CUDA stream, or None when running on CPU.

nvalchemiops.torch.autograd.warp_from_torch(tensor, warp_dtype, requires_grad=None)

Convert a PyTorch tensor to a Warp array with proper gradient tracking.

Parameters:

• tensor (torch.Tensor) – Input PyTorch tensor.

• warp_dtype (wp.dtype) – Warp data type for the resulting array (e.g., wp.vec3f, wp.float64).

• requires_grad (bool, optional) – Override gradient tracking. If None, inherits from tensor.requires_grad.

Returns:

Warp array backed by the same storage as tensor, with gradient tracking enabled when requires_grad is True.

Return type:

wp.array

nvalchemiops.torch.autograd.needs_grad(*tensors)

Check if any of the provided tensors requires gradients.

Useful for conditionally enabling Warp gradient tracking and tape recording only when needed for backpropagation.

Parameters:

*tensors (torch.Tensor) – Tensors to inspect; non-tensor values in tensors are silently skipped.

Returns:

True if at least one tensor has requires_grad=True, False otherwise.

Return type:

bool

Examples

>>> positions = torch.randn(100, 3, requires_grad=True)
>>> charges = torch.randn(100, requires_grad=False)
>>> needs_grad(positions, charges)
True
>>> needs_grad(charges)
False

Autograd Context Manager

class nvalchemiops.torch.autograd.WarpAutogradContextManager(enable)

Conditionally create a Warp tape as a context manager.

Yields a wp.Tape when enable=True so Warp kernel launches are recorded for differentiation. When enable=False a no-op context is used, incurring zero overhead.

Parameters:

enable (bool) – Whether to create a tape for gradient recording.

Yields:

wp.Tape or None – Active tape for recording when enable=True; None otherwise.

Examples

>>> needs_grad_flag = needs_grad(positions, charges)
>>> with WarpAutogradContextManager(needs_grad_flag) as tape:
...     wp.launch(kernel, ...)
>>> if needs_grad_flag:
...     # tape is a wp.Tape instance
...     tape.backward()

nvalchemiops.torch.autograd.attach_for_backward(output, tape=None, **warp_arrays)

Attach Warp tape and arrays to a PyTorch tensor for later retrieval in backward.

Stores tape and each entry in warp_arrays as attributes on output, so that retrieve_for_backward() can recover them inside a custom-op backward pass.

Parameters:

• output (torch.Tensor) – PyTorch tensor to attach attributes to (typically the first output of the forward function).

• tape (wp.Tape, optional) – Warp tape containing recorded operations for the backward pass.

• **warp_arrays (wp.array) – Named Warp arrays to store (e.g., positions=wp_positions, charges=wp_charges).

Return type:

None

Examples

>>> attach_for_backward(
...     output,
...     tape=tape,
...     positions=wp_positions,
...     charges=wp_charges,
...     energies=wp_energies,
... )
>>> # Later in backward:
>>> tape = output._warp_tape
>>> wp_positions = output._wp_positions

nvalchemiops.torch.autograd.retrieve_for_backward(output, *array_names)

Retrieve Warp tape and arrays from a PyTorch tensor in backward pass.

Companion to attach_for_backward(). Arrays not present on output are silently omitted from the returned dict so optional attachments do not need special handling.

Parameters:

• output (torch.Tensor) – PyTorch tensor that has attached Warp objects via attach_for_backward().

• *array_names (str) – Names of Warp arrays to retrieve (without the _wp_ attribute prefix).

Returns:

• wp.Tape – The stored Warp tape attached to output.

• dict[str, wp.array] – Mapping from each requested name to its stored Warp array.

Return type:

tuple[Tape, dict[str, array]]

Examples

>>> tape, arrays = retrieve_for_backward(
...     ctx.output,
...     'positions', 'charges', 'energies'
... )
>>> wp_positions = arrays['positions']
>>> tape.backward()

nvalchemiops.torch.autograd.extract_gradients(ctx, warp_arrays, input_names)

Extract gradients from Warp arrays and return them in PyTorch autograd order.

Returns gradients aligned to the forward function’s input signature: each position receives a torch.Tensor when the corresponding input required gradients, or None otherwise.

Parameters:

• ctx (Any) – PyTorch autograd context. Must expose each name in input_names as an attribute holding the saved input tensor.

• warp_arrays (dict[str, wp.array]) – Mapping from input names to Warp arrays whose .grad fields have been populated by tape.backward().

• input_names (list[str] or tuple[str]) – Names of inputs in the same order they appear in the forward function signature.

Returns:

Gradients in input_names order; None for inputs that did not require gradients or for which no Warp array was provided.

Return type:

tuple[torch.Tensor or None, …]

Examples

>>> # In backward function:
>>> tape, arrays = retrieve_for_backward(ctx.output, 'positions', 'charges')
>>> tape.backward()
>>> return extract_gradients(
...     ctx,
...     arrays,
...     ['positions', 'charges', 'cell', 'alpha']
... )
>>> # Returns: (grad_pos, grad_charges, None, None)

nvalchemiops.torch.autograd.standard_backward(ctx, grad_outputs, output_names, array_names, input_names, output_dtypes=None)

Run the standard Warp tape backward pass for a PyTorch custom operator.

Handles both single-output and multi-output operators by encapsulating the common backward pattern: retrieve tape and Warp arrays from context, copy upstream gradients into the output arrays, call tape.backward(), then return input gradients via extract_gradients().

Parameters:

• ctx (Any) – PyTorch autograd context with saved tensors

• grad_outputs (torch.Tensor or tuple[Optional[torch.Tensor], ...]) – Gradient(s) from upstream operations. - Single output: pass the gradient tensor directly - Multiple outputs: pass tuple of gradient tensors (None if unused in loss)

• output_names (str or Sequence[str]) – Name(s) of the output array(s) stored in ctx. - Single output: ‘output’ or ‘energies’ - Multiple outputs: [‘energies’, ‘forces’]

• array_names (Sequence[str]) – Names of ALL warp arrays that were attached (outputs + inputs). MUST include all output array names first! Examples: - Single output: [‘output’, ‘positions’, ‘charges’] - Multiple outputs: [‘energies’, ‘forces’, ‘positions’]

• input_names (Sequence[str]) – Names of all inputs in forward function signature order

• output_dtypes (Any or Sequence[Any], optional) – Warp dtype(s) for each output. Required for multiple outputs or non-float32 outputs. - Single output: wp.float32 (default) or wp.vec3f - Multiple outputs: [wp.float32, wp.vec3f]

Returns:

Gradients for all inputs (None for those without requires_grad)

Return type:

tuple[Optional[torch.Tensor], …]

Examples

Single output operator:

>>> # In forward:
>>> attach_for_backward(output, tape=tape, output=wp_output,
...                     positions=wp_positions, charges=wp_charges)
>>>
>>> # In backward:
>>> def backward(ctx, grad_output):
...     return standard_backward(
...         ctx,
...         grad_outputs=grad_output,  # Single tensor (note: parameter name)
...         output_names='output',  # Single string
...         array_names=['output', 'positions', 'charges'],
...         input_names=['positions', 'charges', 'cell', 'alpha'],
...     )

Multiple output operator:

>>> # In forward:
>>> attach_for_backward(energies, tape=tape, energies=wp_energies,
...                     forces=wp_forces, positions=wp_positions)
>>> return energies, forces
>>>
>>> # In backward:
>>> def backward(ctx, grad_energies, grad_forces):
...     return standard_backward(
...         ctx,
...         grad_outputs=(grad_energies, grad_forces),  # Tuple
...         output_names=['energies', 'forces'],  # List
...         output_dtypes=[wp.float32, wp.vec3f],  # Required!
...         array_names=['energies', 'forces', 'positions'],
...         input_names=['positions'],
...     )