Pulser icon indicating copy to clipboard operation
Pulser copied to clipboard

Support differentiability through Torch tensors

Open HGSilveri opened this issue 7 months ago • 0 comments

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:

  1. 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 the torch support, they can have it by adding [torch] to their pip install commands:
# Both options work
pip install pulser[torch]
pip install pulser-core[torch]
  1. To handle both np.ndarrays and torch.Tensors indiscriminately, the wrapper class AbstractArray is created. This class contains an array-like object, which is stored either as torch.Tensor or a np.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.
  2. Operations on AbstractArray must use custom functions to preserve the AbstractArray type in the output. As such, analgous definitions of the necessary numpy/torch functions were defined in the new pulser.math module. This new module is, by convention, imported with import pulser.math as pm and can often be used as a drop-in replacement for np.
  3. Parameters supporting differentiability are marked with the new pm.Differentiable type-hint. When this is present, np.ndarray and torch.Tensor (with requires_grad=True or not) can be provided interchangeably. Whenever a parameter is not marked as pm.Differentiable, providing a torch.tensor is not guaranteed to work as intended.

HGSilveri avatar Jul 01 '24 16:07 HGSilveri