Implement Dynamical CPU/GPU/TPU Backends

[Sampreet]

Summary

With reference to the email conversation with @piperfw and @gefux, I am opening a dedicated issue pertaining to a dynamical numerical backend to support GPUs and TPUs. OQuPy utilizes TensorNetwork modules which support frameworks like TensorFlow, PyTorch and JAX. However, TensorNetwork's Node requires its tensor parameter to be passed with the same numerical backend, and OQuPy's modules use NumPy explicitly to create these nodes. As such, using a dynamical backend to switch from NumPy/SciPy to corresponding libraries of (say) JAX will facilitate the usage of TensorNetwork's GPU/TPU-friendly frameworks, which may speed up certain methods. This issue proposes a way similar to what has been recently implemented in QuTiP to support auto-differentiation using its JAX backend.

Implementation

The oqupy.config module is updated with:

# numerical backend
import numpy as default_np
import scipy.linalg as default_la
NUMERICAL_BACKEND_NUMPY = default_np
NUMERICAL_BACKEND_LINALG = default_la
NumPyDtypeComplex = default_np.complex128   # earlier NpDtype
NumPyDtypeFloat = default_np.float64        # earlier NpDtypeReal

The specific choice of scipy.linalg instead of scipy is because jax.scipy.integrate doesn't yet have implementations of dblquad, quad and quad_vec as required by some of OQuPy's modules.
A new oqupy.backends.numerical_backend module now takes care of the switching. The coarse contents of the module are:

class NumPy:
    """
    The NumPy backend for dynamic
    switching through `oqupy.config`.
    """
    @property
    def backend(self) -> default_numpy:
        """Getter for the backend."""
        return oqupy.config.NUMERICAL_BACKEND_NUMPY
    def __getattr__(self, name):
        """Return the backend's default attribute."""
        backend = object.__getattribute__(self, 'backend')
        return getattr(backend, name)
    # additional overridden methods for dataypes,
    # array updates, random number generators, etc

class LinAlg:
    # same as above for ``default_scipy.linalg``

# initialize for import
np = Numpy()
la = LinAlg()

All the modules of OQuPy are modified as:

# remove all `numpy` and `scipy.linalg` dependences other than
# `numpy.ndarray` before importing the dynamical backend
from oqupy.backends.numerical_backend import np, la
# change occurrences of NpDtype and NpDtypeReal
# to the corresponding methods of the new backend

The corresponding changes can be viewed by comparing the pr/feature-numerical-backend branch.

Testing

I have tested the above implementations using tox for vanilla NumPy/SciPy backend and reproduced the plots of arXiv:2406.16650. To use the JAX backend, the following snippet can be added at any point in the scripts:

import jax
import jax.numpy as jnp
import jax.scipy.linalg as jla
import os
import oqupy.config as oc
import tensornetwork as tn
jax.config.update('jax_enable_x64', True)
oc.NUMERICAL_BACKEND_NUMPY = jnp
oc.NumPyDtypeComplex = jnp.complex128
oc.NumPyDtypeFloat = jnp.float64
oc.NUMERICAL_BACKEND_LINALG = jla
tn.set_default_backend('jax')

Issues

The JAX backend tests have multiple issues that require suggestions before any modification. A few of these issues are mentioned below:

Immutability of jax.numpy Arrays

The add_singleton function in oqupy.utils is often called by modules like oqupy.backends.pt_tempo_backend to create/update MPOs and MPSs. This function updates the shape of the array by adding an additional dimension to it at a specified index. Whereas NumPy supports altering the shape attribute, JAX's ArrayImpl datatypes are immutable and as such shape doesn't have a setter method leading to a recurrent errors. Similar mutations can be found in oqupy.gradient, oqupy.mps_mpo and oqupy.system_dynamics modules. In this case, would reshaping the array instead of updating its shape break the functionality? Otherwise, a separate variable can be dedicated for the shape.

Vectorization Signatures

Modules such as oqupy.bath_correlations and oqupy.system use the vectorize method of NumPy without explicitly mentioning the signatures for the input and output. This raises issues with JAX when the parameters contain non-scalar inputs. Depending on the type of implementation, can the signatures be added for each vectorized function? For example, the Hamiltonians for TimeDependentSystem and TimeDependentSystemWithField will have the signatures ()->(m,m) and (),()->(m,m) respectively (kindly correct me if I have made a mistake here).

Row Degeneracy Checks

Unlike the numpy.unique method, the jax.numpy.unique method returns array indices with the same shape as the input array even with the axis parameter. Although this can be resolved by using the flatten method, I am not sure if it will introduce additional errors.

Kindly share your views as and when time permits. Thank you once again.

[piperfw]

Hi @Sampreet, the code changes look really good. I haven't been able to run things on my system yet, as Linux + AMD for JAX is experimental, and there are some barriers to building from source, but hopefully I'll be able to get that done soon.

[On a related note it looks like JAX requires Python 3.10+; not an issue but will need to make a note of it in the README or similar as I think we are still hoping to support 3.6 as a minimum - #80]

The dynamical switching looks fine, and I actually prefer the neatness of using np.function calls instead of importing a load of specific functions in each module. I wonder if we should adjust the config setup slightly, at least e.g. to check DTYPE_COMPLEX and DTYPE_REAL are consistent with BACKEND_NUMPY, or even remove the freedom of being able to set the datatypes - @gefux what do you think? Similarly, if setting the default tensor network backend to 'jax' is important, maybe we want to check or take care of that.

I guess it's nice to be able to choose JAX for either a numpy of scipy.lianlg replacement, or both, and we can pin an issue to review the scipy import if the integrate functions are brought to JAX.

EDIT: Running CPU JAX, I've observed that the calls jax.config.update and tn.set_default_backend are necessary to avoid errors/warnings. On the other hand, JAX seems to understand np.complex128, np.float64 (although we may still want to advise their setting - depending on what Gerald intended for those config variables).

[piperfw]

RE Potential Issues

Shape

I don't see any issue with this. In fact, the current use of shape vs reshape in the code looks sporadic. I think we should switch to always using .reshape, which is safer (reshaping an array in-place will fail if a copy is required). See the deprecation warning https://numpy.org/doc/stable/reference/generated/numpy.ndarray.shape.html

Vectorization Signatures

Yes, signatures can be added and that sounds like a good addition. You are correct: Hamilonians for all System, TimeDependentSystem, TimeDependentSystemWithField must all return (m,m) where m is the Hilbert space dimension (._dimension), and take 0, 1 and 2 scalar parameters respectively (for the mean-field system, it is complex). See e.g. the parsing functions lines 986, 998 of system.py on your branch. I suggest you go ahead and add the signatures as you see fit, if there are others you are unsure on just ask or share a list and Gerald or I should be able to make those changes.

Row Degeneracy Checks

.flatten() there looks absolutely fine. Note that routine is only run a single time at the start of the computation, so there is no worry about overhead of unnecessary copies etc.

Overall, this looks all looks very good and close to being PR-ready. It would be good to do widespread testing, because presumably there may well be other small compatibility issues such as you have caught in the numerical backend with update. That being said, I think we can already (or soon) be confident it does not break the existing/default numpy backend, and can always advertise the JAX backend as more experimental or cutting edge.

[piperfw]

RE Unit testing

We have a large number of coverage and 'physical' (assessing physical consistency or accuracy) under tests/. What is the best way to extend these to test against the JAX backend? I guess (assuming we want the code to work 100% with both backends) we can have two testing environments, one for each backend (JAX also being Python 3.10), to run all the tests against.

[Sampreet]

Hi @piperfw, thanks again for your detailed comments. I shall implement the suggestions for shape, signatures and the degeneracy arrays and let you know if I face any issue. Based on my understanding of the source code, I would like to share the following in relation to your queries:

Data types

I do agree with certain checks on datatypes, specially since the precision requirements (single/double floating points and their corresponding complex datatypes) are interrelated in most of the modules. Also, as the operations for a single simulation (by a user) are likely to involve the same degree of precision, I feel it might not be useful to provide the dynamic alterability of precision (to the user) through the corresponding methods of numerical_backend once the NumPy and LinAlg classes are initialized. Kindly let me know your thoughts on the same. Likewise, I am also interested in knowing the intended purpose of the datatype variables in oqupy.config.

And yes, JAX, by default, supports 32-bit floats (and correspondingly, 64-bit complex) datatypes and requires a call to the jax.config.update method to utilize 64-bit floats. Also, using JAX with TensorNetwork requires a call to the tn.set_default_backend method.

Unit tests

I am still to update the test files with the classes from numerical_backend. As such, I have only tested the branch with vanilla NumPy. To run the tests on the different backends, can we utilize the __init__ files in tests to automatically detect the presence of JAX (with/without CUDA) and set up the backend configurations before the actual tests are run? I guess if those pass for JAX, they will pass for NumPy as well. To be sure about that, we can have two testing environments as you have suggested.

Regarding PR

It might be handy to have a separate branch (say dev-jax) in the main repository for JAX-related changes. In that case, once the necessary changes are done and the tests pass and everything else is finalized, a draft PR can be created to that branch and we can continue with code-specific changes in that draft.

Update on speedup

In the latest commit of the pr/feature-numerical-backend branch, I have implemented the generalized method to create deltas (with a NumPy-friendly oqupy.util.create_delta) which is discussed elaborately in #140. This works seamlessly with JAX and the update introduces a good speedup (benchmarks added inside the issue) for computation and application of PT-MPOs. Kindly share your comments on the same as and when time permits.

About scipy.integrate

The integrate module will probably not be incorporated in the main branch of JAX due to the presence of loopy algorithms. Although, there are probably other libraries that wrap scipy.integrate for JAX. The link might also be useful to plan future development.

Update 2024-09-04

I have made the changes corresponding to shape -> reshape, added vectorization signatures and flattened degeneracy check arrays in the pr/feature-numerical-backend branch, with some other minor changes (list -> ndarray, in-place updates, parameter type-casts and code-style-consistency fixes). I have also added the numerical backend to the test files in tests/coverage and tests/physics. The following additional changes are made in the tests:

• np.testing.assert_almost_equal is changed to default_np.testing.assert_array_almost_equal since jax.numpy doesn't have a testing module and JAX-arrays don't go along without the _array addition. Same goes for np.testing.assert_equal.
• type(<ndarray_element>) is changed to <ndarray_element>.dtype.

All tests pass with both NumPy/SciPy and their corresponding JAX-CPU backends (although slower). The configurations for JAX are added in the corresponding __init__ files (commented out) of tests/coverage and tests/physics. I would need suggestions regarding the following:

• isinstance(<vectorized_function>, np.vectorize) in tests/coverage/system_test doesn't work with jax.numpy.vectorize. I'm not sure if there is a workaround for that. I have commented out the assertions pertaining to those for now.
• the set_initial_tensor method of oqupy.process_tensor.SimpleProcessTensor throws a future warning regarding the creation of a JAX-array with the element None is treated as NaN. I think the tensor element can be set to NaN to avoid the warning or handled separately while calling the initial tensor.

Update 2024-09-05

The JAX-GPU backend doesn't yet support eigen decomposition for non-symmetric matrices although numpy.linalg lies in the scope for JAX. This leads to issue with initialization of oqupy.bath.Bath for non-diagonal coupling operators. Since it is a single operation, using the eigen decomposition in the CPU backend and passing the eigenvectors back to the GPU is a workaround. Otherwise, contributing libraries like jaxeigs can be used. A curated list of JAX-based projects can be found in awesome-jax. With the workaround mentioned, all tests pass on the JAX-GPU backend (taking a good amount of time) without any additional errors/warnings.

Thank you.

[piperfw]

Hi Sampreet, thanks for the detailed updates. I will respond to particular comments below.

The code changes and generally the implementation look really good. For the bigger picture, there are two things we want to avoid:

1. increasing the maintenance difficultly of the codebase and
2. increasing barriers for new contributions

These are potentially conflicting: if we require contributions to be JAX compatible, then that is a large barrier, but if we don't then contributed code may break for the JAX backend and someone will need to assess and fix any issues.

It seems like we should be able to keep the core OQuPy and existing features JAX-compatible (to the extent of lacking integrate, eigen decomposition etc.). I suggest we add a note to developers (and forgetful maintainers!) e.g. in CONTRIBUTING.MD to advise on the JAX backend, and best practices. The main ones I can see being to use the dynamical imported np. module rather than explicit np functions, updating array values using the new update method, and passing array-type arguments explicitly as arrays, not lists. Can you help provide any points I'm missing / advisory text?

I'm not sure we should absolutely require developers to write modules for new features that are necessarily JAX compatible. I guess there might be some boilerplate code we can have ready to add to a module e.g. to raise NotImplemented if the JAX backend is loaded for an incompatible module (similarly, one would need to skip tests of that module with the JAX backend). I don't think this point needs any action now.

Otherwise, I realised we did in fact switch to Python3.10 making my concern about versions above irrelevant. So that's good.

Data types

I'll let @gefux tell us about the intended purpose, and if this is still relevant, as I don't know.

Unit tests

I'm currently having difficultly getting a working JAX install on my system within python testing environments (3.10 is not my OS python). I'll hopefully get this sorted in due course and provide an update. But regardless, the assertion and vectorized changes sound fine. I don't know of any code that would suffer from the change of None to NaN in SimpleProcessTensor, but I'll highlight this for @gefux in case he does.

My question here Sampreet is how we can make switching to the JAX testing neater (bear in mind we want to use GitHub's CI). I'm not so familiar with pytest, but do you think there is a simple command flag or plugin we can optionally specify from the test invocation command to optionally load the JAX __init__.py code that is currently commented out?

PR

That's a good suggestion. For now I think just creating a PR from your current branch and linking here would be helpful so everyone can follow changes.

create_delta

I or someone else will look into #140 when we get the chance, sounds good.

integrate / linalg

These deficiencies may be good to note in addition to advise for developers. Maybe we need a page in the docs about the JAX backend, what do you think? Under 'Development', perhaps a new page on JAX after 'Contributing' (https://oqupy.readthedocs.io/en/latest/pages/contributing.html). In place of or addition to changes to CONTRIBUTING.md suggested above.

I do not think we should use other libraries or projects. In particular, the bath eigendecomposition is only done once for small matrices if I'm not mistaken, so there is no cost in just avoiding the GPU here.

[gefux]

Dear @Sampreet,

I finally found some time to look into your code - I apologize for the delay.
Overall, I agree with @piperfw. I think the changes look good.

General comment

As piperfw writes, we need to be mindful of the maintainance difficulty and the contributing barriers. I think that using from oqupy.backends.numerical_backend import np instead of import numpy as np is completely tollerable from the point of view of potential future contributors. On the other hand, demanding a JAX compatability for all current and future modules could be a discouraging factor for many potential contributors, or at least then leave more work for the maintainers that would then have to make the code JAX compatible. I'd thus suggest to make the changes you propose, but keeping the focus of OQuPy on vanilla numpy and declaring the JAX option as 'experimental' for users and as 'optional' for contributors (at least for now).

Technical comments (in addition to piperfw's)

Dynamical switching and automated testing:

The general idea of defining the default in config.py and then constructing the np and sp modules in backends/numerical_backend.py looks good to me. The switching, however, seems to be done by dynamically changing the variables defined in config.py. I think that the variable default_np should not be touched when switching, but instead the backends/numerical_backend.py should check a system environment variable such as "OQUPY_BACKEND". If it is set to "JAX" then np should be set to be the jax module and otherwise np should be set to the default. As JAX litterally requires a different system environment I think that a system environment variable is the right place to signal the switch. This would eliminate the __init__.py in the testing directory. The automated tests would then be done in a separate "JAX" environment that has everything neccessary installed and the "OQUPY_BACKEND=JAX" variable set. For functionality that doesn't work on JAX one could then add @skip decorators that are conditioned on the backend type.

Numpy data types

Fixing the data types to one specific type is simply a kind of a defensive coding style which I embrased when starting this project because numpy tends to become really slow if it performs computations on mixed types. At least that is what I used to experience a few years ago -- maybe this has changed in the meantime. I also found it potentially useful to be able to switch all computations to a different type from one point.

None vs NaN in SimpleProcessTensor.set_initial_tensor

If I am not mistaken it looks to me that the line self._initial_tensor = None there is completely useless as it is overwritten in the next line anyways. In the SimpleProcessTensor.__init__() the line self._initial_tensor = None however is important for the logic of the class as it signals that the process tensor doesn't have an initial tensor (which is still a valid form of a process tensor). It would be important that we could keep this logic of a tensor being either None or a tensor depending on it's existence because this appears in many places in the code, and is very convenient.

[piperfw]

Really good idea with the environmental variable @gefux, that clears up a lot of the mess with the testing I was concerned about. Completely agree about the experimental label for the backend; we can advertise this in an appropriate page in the docs.

[Sampreet]

Dear @piperfw and @gefux,

Thanks a lot for the detailed comments. I shall implement the suggested changes sometime this week and create a draft PR, where we can continue the code-related discussions. Below are a few points I would like to add here:

Contribution and Maintenance

It seems like we should be able to keep the core OQuPy and existing features JAX-compatible (to the extent of lacking integrate, eigen decomposition etc.) - @piperfw
I'd thus suggest to make the changes you propose, but keeping the focus of OQuPy on vanilla numpy and declaring the JAX option as 'experimental' for users and as 'optional' for contributors (at least for now). - @gefux

This sounds like a good idea and more user/contributor/maintainer-friendly for now. The proposed changes are also fully compatible with vanilla NumPy users without requiring any change in the existing examples. Same goes for contributions with vanilla NumPy. Additionally, interested contributors can follow an "optional" set of guidelines (detailed next) to make sure that their changes work seamlessly with JAX, thereby easing maintenance. Moreover, the current oqupy.backends.numerical_backend module simply handles the breaking changes in JAX NumPy such as in-place updates and pseudo random number generation (which defaults to vanilla NumPy for now), and therefore, no explicit JAX-related imports are required in any of the modules. This keeps JAX-dependencies away from the core library, and the users can utilize the oqupy.config module to switch to JAX-based "experimental" features dynamically. Alternatively, a new function (say enable_experimental_features) can be added to ease the switching but will require explicit JAX-imports. I shall detail on this a bit more in the section for environment variables.

I suggest we add a note to developers (and forgetful maintainers!) e.g. in CONTRIBUTING.md to advise on the JAX backend, and best practices. Can you help provide any points I'm missing / advisory text? - @piperfw

The following guidelines can be added for "optional" contribution to the "experimental" features :

• using from oqupy.backends.numerical_backend import np instead of import numpy as np and use the alias default_np for vanilla NumPy; avoiding wildcard imports.
• using from oqupy.backends.numerical_backend import la instead of import scipy.linalg as la; except for non-symmetric eigen-decomposition (scipy.linalg.eig for non-symmetric matrices).
• using one of np.dtype_complex (np.dtype_float) or oqupy.config.NumPyDtypeComplex (oqupy.config.NumPyDtypeFloat) instead of np.complex_ (np.float_).
• converting lists or tuples to NumPy arrays when passing them as arguments inside NumPy functions.
• using np.update(array, indices, values) instead of array[indices] = values.
• using np.get_random_floats(seed, shape) instead of np.random.default_rng(seed).random(shape).
• declaring signatures for np.vectorize explicitly.

Maybe we need a page in the docs about the JAX backend, what do you think? Under 'Development', perhaps a new page on JAX after 'Contributing'. In place of or addition to changes to CONTRIBUTING.md suggested above. - @piperfw

It might be a good idea to add a separate section/page in the documentation to highlight the "experimental" nature of this feature. In time, this can include installation guidelines, gotchas (some points from CONTRIBUTING.md), and configurational changes required to use JAX with OQuPy.

Continuous Integration and Unit Tests

My question here Sampreet is how we can make switching to the JAX testing neater (bear in mind we want to use GitHub's CI). - @piperfw

From what I understand, the automated tests are performed with tox using the GitHub CI routine described in .github\workflows\python-package-tests.yml. For the current changes, we can therefore add a separate testenv block in tox.ini:

[testenv:pytest_jax]
description = run pytests with JAX and make coverage report
basepython = python3.10
deps = -rrequirements_ci_jax.txt
commands =
    pytest --cov-report term-missing --cov=oqupy ./tests/{posargs}
    python -m coverage xml

Then, we can use Python's importlib to run the tests with and without JAX via the __init__.py modules of tests/coverage and tests/physics:

from importlib.util import find_spec

if find_spec('jax') is not None:
    # JAX configuration here

I have implemented this approach and it works well. What I am worried about is the time taken by the JAX-based tests are very high since we do not have any explicit JAX-based primitives in OQuPy's modules yet (further details in the next section).

Dynamical Switching with Environment Variables

I think that the variable default_np should not be touched when switching, but instead the backends/numerical_backend.py should check a system environment variable such as "OQUPY_BACKEND". If it is set to "JAX" then np should be set to be the jax module and otherwise np should be set to the default. As JAX literally requires a different system environment I think that a system environment variable is the right place to signal the switch. This would eliminate the __init__.py in the testing directory. - @gefux

I believe what you are suggesting here is an explicit backend module for JAX, which is imported during the initialization of numerical_backend.NumPy class only when the environment variable "OQUPY_BACKEND" is set to "JAX". Otherwise, the initialization defaults to the NumPy library defined by NUMERICAL_BACKEND_NUMPY in oqupy.config (which stays untouched). Kindly let me know if I am mistaken here.

The above description is similar to what has been implemented inside tensornetwork. I have also implemented something similar for my RL library quantrl by using an abstract base class inherited by each numerical backend (NumPy/JAX/PyTorch), with individual overridden abstract methods. This facilitates a fine-grained control (and thereby, speedups) over each backend-based implementation (such as precision, in-built methods or even custom wrappers for functions). The idea of utilizing oqupy.config here instead of this approach was to initially support GPU/TPU arrays without major changes in the original codebase (e.g. syntax of conditional statements, loops will be governed by the backend methods) and avoiding any explicit mention for JAX inside it.

On second thoughts, since tensornetwork is already a required module for OQuPy, we can utilize its JAX backend for this purpose and wrap it via numerical_backend based on the value of "OQUPY_BACKEND" (a few additional changes might make it compatible with PyTorch/Symmetric/TensorFlow too). This will allow:

• the use of JIT-ted functions and jax.lax primitives inside the modules, resulting in speedups in computation and application.
• testing without the requirement of explicit JAX configurations (setting "OQUPY_BACKEND" would still require importlib as mentioned in the previous section).
• changing the backend dynamically by using methods like enable_experimental_features and disable_experimental_features (this will require reinitialization of numerical_backend.NumPy).
• unavailable JAX functions to be automatically overridden by vanilla NumPy ones though the numerical backend.

As of now, I don't see any major disadvantage in the environment-variable approach, other than the fact that we might need to add a few additional guidelines for "optional" contributions.

[piperfw]

Hi @Sampreet, thanks for making the PR and apologies for the delay in responding here.

Contribution and Maintenance

I do think we should keep to no required jax imports in modules, but something like enable_experimental_features sounds fine. The optional guidelines you list is perfect, and I agree with a separate page on the docs for this - it doesn't have to be lengthy.

Continuous Integration and Unit Tests

Your proposal for tox.ini also looks good, thanks for implementing this approach. If test runtime is going to be problematic for GitHub CI, then we can use @Skip as Gerald suggests or omit jax-testing when building via GitHub CI entirely. Howeer, I think it is useful to at least have the ability to run the tests (locally) with the jax backend.

Dynamical Switching with Environment Variables

I understand @gefux's suggestion to mean, the dynamical switching in numerical_backend.py is the same as you have designed, only instead of switching based on settings in config.py you switch base on an environment variable "OQUPY_BACKEND" (which if unset should result in numpy). Then one can have support for GPUs without major changes in the codebase as you comment was the original motivation. This sounds different to what you suggest, an explicit jax backend module? Please let me know if I am misunderstanding.

Being able to use jax primitives etc. does sound really nice; that would really open up possibilities for developers to push and optimise oqupy (/for GPUs), although I'm not sure about the code complexity.

I'll make a brief comment/summary on the PR itself.

[Sampreet]

Hi @piperfw, thank you for your the feedback and the suggestions for the next steps. In line with your comment on Contribution and Maintenance, I shall update the docs sometime during the weekend. As for CI and Unit Tests, I think it's best if the JAX-based tests are run locally since GitHub CI has a monthly capping for it's free tier. To achieve this using enable_experimental_features, the numerical_backend.py file can be modified as (note that pylint throws 'import-outside-toplevel' for the try-catch block):

# store instances of the initialized backends
# this way, `oqupy.config` remains unchanged
# and `OQUPY_BACKEND` is used by default
# when NumPy and LinAlg are initialized
NUMERICAL_BACKEND_INSTANCES = {}

def get_numerical_backends(
        backend_name:str,
    ):
    if backend_name in NUMERICAL_BACKEND_INSTANCES:
        return NUMERICAL_BACKEND_INSTANCES[backend_name]
    assert backend_name in ['jax', 'numpy'], \
        "currently supported backends are `'jax'` and `'numpy'`"
    if 'jax' in backend_name.lower():
        try:
            # explicitly import and configure jax 
            # tox catches import-outside-toplevel
            import jax
            import jax.numpy as jnp
            import jax.scipy.linalg as jla
            jax.config.update('jax_enable_x64', True)
            # set tensor network backend
            set_default_backend('jax')
            NUMERICAL_BACKEND_INSTANCES['jax'] = [jnp, jla]
            return NUMERICAL_BACKEND_INSTANCES['jax']
        except ImportError:
            print("JAX not installed, falling back to NumPy")

    set_default_backend('numpy')
    NUMERICAL_BACKEND_INSTANCES['numpy'] = [default_np, default_la]
    return NUMERICAL_BACKEND_INSTANCES['numpy']

# NumPy and LinAlg initializes with the `OQUPY_BACKEND` value by default
class NumPy:
    def __init__(self):
        self.backend = get_numerical_backends(oc.OQUPY_BACKEND)[0]
    # other methods remain the same
class LinAlg:
    def __init__(self):
        self.backend = get_numerical_backends(oc.OQUPY_BACKEND)[1]
    # other methods remain the same

# initialize
np = NumPy()
la = LinAlg()

# switch the backend without changing `oqupy.config`
def enable_experimental_features():
    backends = get_numerical_backends('jax')
    np.backend = backends[0]
    la.backend = backends[1]

I suppose this is in line with @gefux's comment. Kindly let me know if this works.

Also, to answer your query on explicit modules, what I meant was that the alternate form of implementation which I mentioned in my previous comment would use separate modules for each numerical backend (as in quantrl), and implement an abstract base class inherited by the numerical backends. This approach typically follows the file structure:

oqupy/
└───backends/
    ├───__init__.py
    ├───base.py
    ├───jax.py
    ├───numpy.py
    └───torch.py

Here, each backend can be individually optimized to implement the required methods. Although this will enhance scalability and allow further optimization of the package (as you have mentioned), the codebase might change significantly. As such, this can be implemented at a later stage when OQuPy's backends are modified to implement other planned changes (as mentioned in some of the ToDos and in #28).

Thanks again!

[piperfw]

Hi @Sampreet! I see what you're doing here and this looks close to what we want. Thanks for explaining the explicit module idea, I agree we can leave that for the time being.

To check, is the purpose of NUMERICAL_BACKEND_INSTANCES to prevent all the 'numpy intialisation' (which is really just tn.set_default_backend('numpy')) being run twice when one initalises the np and la instances? Or to allow one to 'hot switch' back and forth between numpy and jax after using the other (I'm not sure that would work as you would need to recall tn.set_default_backend at least)?

I think the main idea of the dynamical switching was to switch based on the environment, which I think can be done with something like

# config.py
OQUPY_BACKEND_ENV_VAR = 'OQUPY_BACKEND' # name of environmental variable to check

# numerical_backend.py
import os

def get_numerical_backends(...)

class NumPy:
    def __init__(self, backend_name=oc.OQUPY_DEFAULT_BACKEND):
        self.backend = get_numerical_backends(backend_name)[0]
# Similar for LinAlg


try:
    backend_name = os.environ[oc.OQUPY_BACKEND_ENV_VAR]
except KeyError:
    backend_name = oc.OQUPY_DEFAULT_BACKEND

np = NumPy(backend_name=backend_name)
la = LinAlg(backend_name=backend_name)

Note I think it makes sense to call the config variable OQUPY_DEFAULT_BACKEND rather than just OQUPY_BACKEND as it does not necessarily reflect the backend in use (if the environmental variable was set).

In fact, can one lose NUMERICAL_BACKEND_INSTANCES entirely if the Numpy, LinAlg classes take a backend directly?

class NumPy:
    def __init__(self, backend=None):
        if backend is None:
            backend = get_numerical_backends(oc.OQUPY_DEFAULT_BACKEND)
        self.backend = backend

# Similar for LinAlg
np_backend, la_backend = get_numerical_backends(backend_name)
np = NumPy(np_backend)
la = LinAlg(la_backend)

I guess that might not be advisable because someone could try to initialise with modules that was not canonically jax or numpy.

What do you think? To summarise, I think this is close to being what we want but I am questioning the relevance of NUMERICAL_BACKEND_INSTANCES

[Sampreet]

Hi @piperfw,

Indeed, the main idea behind the NUMERICAL_BACKEND_INSTANCES was to prevent the multiple initialization of the same library (and subsequently reduce the assertion and configuration checks). In fact, the get_numerical_backends function does something very similar to what tn.set_default_backend does behind the hood. But as you have pointed out, I missed the tn.set_default_backend statement before it returns the instantiated backend. Thank you. One can also hot-switch between the libraries now using enable_jax_features but that might require some additional type-handling (say, converting tensors loaded from file as NumPy array to a JAX array) on the user end for some cases.

Also, I do agree that allowing NumPy and LinAlg classes to take in the library directly might cause issues if users try to pass libraries other than the integrated ones.

I have made some slight changes to the numerical_backend module to implement the environment variable approach in the new commit. Kindly let me know if that works.

[piperfw]

After some further discussion on the PR page, completed with 6ccbd2e