PropagatorOutput: DataParallel and pytrees

deepwave.PropagatorOutput is deliberately not usable as a full torch.nn.DataParallel return value on more than one GPU. The user-facing consequence is documented on the class itself (return a tensor or a plain tuple of tensors from a forward() that will be wrapped in DataParallel, and see Distributed (multi-GPU) execution); this page records why supporting it was rejected.

One fact decides both halves of this page: the field names belong to the instance, not to the class. That is what makes DataParallel.gather impossible and what the pytree registration has to work around.

Why gathering cannot work

DataParallel.gather reconstructs a module’s return value via type(out)(map(gather_map, zip(*outputs))) – the same convention it uses for a plain namedtuple. PropagatorOutput looks like a namedtuple from the outside (out.name, _asdict()) but cannot support this reconstruction path: unlike a real namedtuple, whose field names are fixed by its class, PropagatorOutput’s names vary per instance (they depend on ndim and, for viscoacoustic, q_n_mechanisms) and are therefore not recoverable from type(out) alone – there is no instance to read _names from at reconstruction time.

Because gather never reaches the point of building an output, info is unreachable under it too, in the same way that it is unreachable under torch.utils.checkpoint(..., use_reentrant=True).

The rejected workaround

Making __new__ accept a single positional iterable of values to paper over this would only ever produce an unnamed tuple (all of the safety this class exists to provide, silently gone) while also making PropagatorOutput(some_iterable) – currently a loud, immediate error – degrade into a silent, wrong construction whenever a caller passes names without values by mistake.

Neither trade-off is worth taking for a DataParallel-only workaround, so this is unsupported by design. __new__ instead raises a DeepwaveTypeError naming DataParallel.gather as the likely cause and pointing at the fix, so the failure is immediate and self-explaining rather than a silently unnamed tuple.

The supported pattern is for a forward() wrapping a propagator to return a single tensor (e.g. out.receiver_amplitudes) or a small tuple of plain tensors instead of a full PropagatorOutput when used with more than one GPU under DataParallel. DistributedDataParallel does not gather this way and is unaffected.

Why the pytree context carries the names, not just info

The same fact that stops DataParallel.gather working decides the shape of the pytree registration. A pytree node splits an object into children (the parts a transform maps over) and a context (everything needed to put it back together). The children are the obvious choice – the tuple elements, in return order, with auxiliary_state expanding into the tensors it holds, so the leaves of an output are exactly its tensors and the flattened order is the documented positional return order.

The context has to hold the field names as well as info. Rebuilding an output means calling PropagatorOutput(names, values, info=info), and the names are not recoverable from the class: they depend on ndim and, on the viscous propagators, on the number of relaxation mechanisms. A registration that carried only info could flatten an output but never unflatten one, so every transform that rebuilds its result – vmap, a has_aux gradient, a tree_map – would fail at the point of return.

info travels with them rather than being dropped or recomputed because a transform does not change how the run was discretised: the internal time step, the resolved PML and the survey region of a vmapped call are the ones the mapped call resolved. The two per-instance caches (final_state and the receiver views) are deliberately not carried: they are derived from the children and the names, so a rebuilt output regenerates them on demand and cannot carry a stale one.

deepwave.AuxiliaryState is registered on the same principle, and for it the context matters even more: it holds the tensor names and every recorded run parameter, including format_version. Those are what a continuation checks itself against, so a state that lost them in a round trip would be accepted by a call it does not belong to, or – worse – have a bundle written to an older layout read as this release’s.

Why a treespec holding one is not serialisable

The same context is what puts JSON serialisation of a tree structure out of reach. torch.utils._pytree.treespec_dumps writes a treespec – the shape of a tree, with each node’s context embedded – as JSON, which torch.export uses. An output’s context holds a deepwave.RunInfo, and an auxiliary state’s holds its recorded run parameters; neither is JSON, and inventing a JSON encoding for them would be inventing a serialisation format for the state of a run, with the compatibility promise that implies, in order to serve a tree shape. So neither type registers a serialised type name, and treespec_dumps on a structure containing one raises PyTorch’s NotImplementedError saying that no serialisation name is registered for it.

Nothing else is affected: flattening, unflattening, key paths, every torch.func transform, torch.compile, tree_map, and torch.save/torch.load or pickle of the objects (which use the ordinary pickle protocol, not a treespec) work as usual. Saving a state to a file is what deepwave.AuxiliaryState’s format_version exists for, and that path is untouched.

What the registration needs from PyTorch

The registration is an enhancement rather than a requirement: a propagator returns and accepts its output either way, and import deepwave must not depend on it. deepwave._pytree therefore probes for each capability it uses and degrades instead of failing:

• torch.utils._pytree.register_pytree_node. Without it nothing is registered and both types stay pytree leaves – an opaque whole, which is what they were before they were registered at all. Transforms that need to see inside them then fail at the point of use, as they did before.

• the flatten_with_keys_fn keyword and the GetAttrKey/MappingKey classes, which arrived later than the function. Without them both types are still registered and every transform still works; only the names of the leaves are lost, so tree_flatten_with_path cannot key them.

A registration already in place – another library that registered the same type, or a module reload that runs the code twice – is left alone rather than overwritten, because objects flattened under it may still be held. None of these cases warns: an older PyTorch losing an optional convenience is not something a user of it can act on.

Why the class still subclasses tuple

That decision is load-bearing for gradient propagation through torch.autograd.Function and torch.utils.checkpoint, and its rationale stays in deepwave.PropagatorOutput’s own docstring rather than moving here, because anyone tempted to change the base class needs to read it at the point of change.