qutip / qutip-jax

JAX backend for QuTiP
BSD 3-Clause "New" or "Revised" License
15 stars 7 forks source link

Allow other `dtype` than complex128. #22

Open Ericgig opened 1 year ago

Ericgig commented 1 year ago

Our data layer always use double precision complex, but this is not always ideals, GPUs often don's support double precision well.

With this, the default is still to convert any array to complex128, but it can be overwritten. However, this is not usable trough the Qutip interface with this PR. I am not too sure how to make it available to the user....

coveralls commented 1 year ago

Pull Request Test Coverage Report for Build 5455627965


Changes Missing Coverage Covered Lines Changed/Added Lines %
src/qutip_jax/qobjevo.py 19 23 82.61%
<!-- Total: 58 62 93.55% -->
Totals Coverage Status
Change from base Build 5378410853: -0.1%
Covered Lines: 742
Relevant Lines: 822

💛 - Coveralls
Ericgig commented 1 year ago
* Should we consider registering a `"jax64"` dtype to addres the complex64 case? In qutip-tensorflow I used a _BaseTfTensor class to represent both types of tensors. I then wrapped this class to create `TfTensor64` and `TfTensor128`. So in qutip-tensorflow these two classes are different but virtually the same. It does the work but there may be cleaner approaches? It does require to register the specialisation twice (or perhaps three times if you want to allow operations between 64 and 128 types).

I would prefer not to. I am interested in using float64 arrays in specific integrators. But these should not be registered since they would not work anywhere else. As it is right now, they would just be converted to complex arrays when needed.

Also having a jax64 type would encourage users to use them, but some have already trouble differentiating floating point error with physics / code error with double precision. It would run into issues with default options (integrator's atol is 1e-8, core options' atol is 1e-12). Many users would not think of updating these settings or understand the error these settings would cause.

* I see the use of fast_constructor instead of the **init** method. Is it to provide a significant speedup in the construction time?