nvalchemi.distributed.autograd_target

nvalchemi.distributed.autograd_target(t)

Return the tensor to pass as torch.autograd.grad()’s inputs=.

Under domain decomposition the framework wraps data.positions (and data.charges) as a ShardTensor view of a halo-padded leaf via Tensor.as_subclass(). The view is not itself in the autograd graph — only the underlying tensor is — so passing the view directly to torch.autograd.grad() raises “differentiated Tensors appears to not have been used in the graph”. This helper returns the in-graph leaf instead.

Call this once where the autograd target is set up (e.g. in adapt_input); the wrapper stays distribution-unaware otherwise.

Parameters:

t (torch.Tensor) – The tensor to differentiate against — a plain tensor, or a ShardTensor view of a halo-padded leaf.

Returns:

The underlying (n_padded, *F) leaf when t is a ShardTensor with a captured autograd source (the halo-padded positions case); t itself otherwise (single-process, where t is already the right target).

Return type:

torch.Tensor