Pulser
Pulser copied to clipboard
Support differentiability through Torch tensors
This PR adds native support for torch.Tensor
on select quantities, so that the final Hamiltonian can be differentiated wrt to these parameters. A few important design considerations were taken here:
-
torch
will be an optional dependency: Users that don't require this feature should be able to keep installing and using pulser as before. For those who do want thetorch
support, they can have it by adding[torch]
to theirpip install
commands:
# Both options work
pip install pulser[torch]
pip install pulser-core[torch]
- To handle both
np.ndarray
s andtorch.Tensor
s indiscriminately, the wrapper classAbstractArray
is created. This class contains an array-like object, which is stored either astorch.Tensor
or anp.ndarray
. To the largest extent possible, instances of this class can be handled like array, allowing operations, indexing, etc... However, their use should be only internal, so they contain only the necessary functionality for the existent codebase - whenever extra functionality is needed, it can always be added. - Operations on
AbstractArray
must use custom functions to preserve theAbstractArray
type in the output. As such, analgous definitions of the necessarynumpy/torch
functions were defined in the newpulser.math
module. This new module is, by convention, imported withimport pulser.math as pm
and can often be used as a drop-in replacement fornp
. - Parameters supporting differentiability are marked with the new
pm.Differentiable
type-hint. When this is present,np.ndarray
andtorch.Tensor
(withrequires_grad=True
or not) can be provided interchangeably. Whenever a parameter is not marked aspm.Differentiable
, providing atorch.tensor
is not guaranteed to work as intended.