Repository: tempoCollaboration/OQuPy

Issue:
Implement Dynamical CPU/GPU/TPU Backends
## 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](https://tensornetwork.readthedocs.io/en/latest/stubs/tensornetwork.Node.html#tensornetwork.Node.__init__), 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](https://github.com/qutip/qutip) to support auto-differentiation using [its JAX backend](https://github.com/qutip/qutip-jax).

## Implementation

The `oqupy.config` module is updated with:

```python
# 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:

```python
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:
```python
# 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](https://github.com/Sampreet/OQuPy/tree/pr/feature-numerical-backend).

## Testing

I have tested the above implementations using `tox` for vanilla NumPy/SciPy backend and reproduced the plots of [arXiv:2406.16650](https://arxiv.org/abs/2406.16650). To use the JAX backend, the following snippet can be added at any point in the scripts:

```python
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](https://jax.readthedocs.io/en/latest/_autosummary/jax.Array.html#jax.Array) 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.

Hints from issue comments:
@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](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision) to utilize 64-bit floats. Also, using JAX with TensorNetwork requires a call to [the `tn.set_default_backend` method](https://github.com/google/TensorNetwork?tab=readme-ov-file#different-backend-support).

### 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](https://github.blog/news-insights/product-news/introducing-draft-pull-requests/) 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](https://github.com/Sampreet/OQuPy/tree/pr/feature-numerical-backend), 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](https://jax.readthedocs.io/en/latest/jep/18137-numpy-scipy-scope.html#scipy-integrate). 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](https://github.com/Sampreet/OQuPy/tree/pr/feature-numerical-backend), 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](https://jax.readthedocs.io/en/latest/jep/18137-numpy-scipy-scope.html#numpy-linalg-numpy-fft). 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](https://github.com/qiyang-ustc/jaxeigs) can be used. A curated list of JAX-based projects can be found in [awesome-jax](https://github.com/n2cholas/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](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html) 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`:

```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`:

```python
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`](https://github.com/google/TensorNetwork/tree/master/tensornetwork/backends). I have also implemented something similar for my RL library [`quantrl`](https://github.com/Sampreet/quantrl/tree/master/quantrl/backends) 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):

```python
# 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](https://github.com/Sampreet/quantrl/tree/master/quantrl/backends)), 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
```python
# config.py
OQUPY_BACKEND_ENV_VAR = 'OQUPY_BACKEND' # name of environmental variable to check
```
```python
# 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?
```python
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](https://github.com/google/TensorNetwork/blob/master/tensornetwork/backends/backend_factory.py). 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

Solution patch:
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 3a3ed40..1c3c7e5 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -29,6 +29,9 @@ The current setup uses:
 * [tox](https://tox.readthedocs.io) ... for testing with different environments.
 * [travis](https://travis-ci.com) ... for continuous integration.
 
+We are actively incorporating additional features to OQuPy,
+details of which can be found in [DEVELOPMENT.md](./DEVELOPMENT.md).
+
 ## How to contribute to the code or documentation
 Please use the
 [Issues](https://github.com/tempoCollaboration/OQuPy/issues) and
diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md
new file mode 100644
index 0000000..ce02df3
--- /dev/null
+++ b/DEVELOPMENT.md
@@ -0,0 +1,43 @@
+# Development
+
+The current development branch "dev/jax" implements
+
+* [Experimental Support for GPUs/TPUs](#experimental-support-for-gpustpus)
+
+## Experimental Support for GPUs/TPUs
+
+Although OQuPy is built on top of the backend-agnostic
+[TensorNetwork](https://github.com/google/TensorNetwork) library,
+OQuPy uses vanilla NumPy and SciPy throughout its implementation.
+
+The "dev/jax" branch adds supports for GPUs/TPUs via the
+[JAX](https://jax.readthedocs.io/en/latest/) library.
+A new `oqupy.backends.numerical_backend.py` module handles the
+[breaking changes in JAX NumPy](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html),
+while the rest of the modules utilizes `numpy` and `scipy.linalg` instances from there
+without explicitly importing JAX-based libraries.
+
+### Enabling Experimental Features
+
+To enable experimental features switch to the `dev/jax` branch and use
+```python
+from oqupy.backends import enable_jax_features
+enable_jax_features()
+```
+Alternatively, the `OQUPY_BACKEND` environmental variable may be set to `jax` to
+initialize the jax backend by default.
+
+### Contributing Guidelines
+
+To contribute features compatible with the JAX backend,
+please adhere to the following set of guidelines:
+
+* avoid wildcard imports of NumPy and SciPy.
+* use `from oqupy.backends.numerical_backend import np` instead of `import numpy as np` and use the alias `default_np` in cases vanilla NumPy is explicitly required.
+* use `from oqupy.backends.numerical_backend import la` instead of `import scipy.linalg as la`, except that for non-symmetric eigen-decomposition, `scipy.linalg.eig` should be used.
+* use one of `np.dtype_complex` (`np.dtype_float`) or `oqupy.config.NumPyDtypeComplex` (`oqupy.config.NumPyDtypeFloat`) instead of `np.complex_` (`np.float_`).
+* convert lists or tuples to arrays when passing them as arguments inside functions.
+* use `array = np.update(array, indices, values)` instead of `array[indices] = values`.
+* use `np.get_random_floats(seed, shape)` instead of `np.random.default_rng(seed).random(shape)`.
+* declare signatures for `np.vectorize` explicitly.
+* avoid directly changing the `shape` attribute of an array (use `.reshape` instead)
diff --git a/docs/index.rst b/docs/index.rst
index 8f60b2a..9b71f4b 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -112,6 +112,7 @@ Furthermore, OQuPy implements methods to ...
    :caption: Development
 
    pages/contributing
+   pages/gpu_features
    pages/authors
    pages/how_to_cite
    pages/sharing
diff --git a/docs/pages/api.rst b/docs/pages/api.rst
index 947b106..1ffb5ca 100644
--- a/docs/pages/api.rst
+++ b/docs/pages/api.rst
@@ -178,7 +178,6 @@ class :class:`oqupy.pt_tebd.PtTebd`
     dictionary.
 
 
-
 Results
 -------
 
@@ -207,3 +206,6 @@ module :mod:`oqupy.operators`
 function :func:`oqupy.helpers.plot_correlations_with_parameters`
   A helper function to plot an auto-correlation function and the sampling
   points given by a set of parameters for a TEMPO computation.
+
+function :func:`oqupy.backends.enable_jax_features`
+  Option to use JAX to support multiple device backends (CPUs/GPUs/TPUs).
diff --git a/docs/pages/gpu_features.rst b/docs/pages/gpu_features.rst
new file mode 100644
index 0000000..5dd2a15
--- /dev/null
+++ b/docs/pages/gpu_features.rst
@@ -0,0 +1,55 @@
+Experimental Support for GPUs/TPUs
+==================================
+The current development branch "dev/jax" implements experimental support
+for GPUs/TPUs.
+
+Although OQuPy is built on top of the backend-agnostic
+`TensorNetwork <https://github.com/google/TensorNetwork>`__ library,
+OQuPy uses vanilla NumPy and SciPy throughout its implementation.
+
+The "dev/jax" branch adds supports for GPUs/TPUs via the
+`JAX <https://jax.readthedocs.io/en/latest/>`__ library. A new
+``oqupy.backends.numerical_backend.py`` module handles the
+`breaking changes in JAX
+NumPy <https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html>`__,
+while the rest of the modules utilizes ``numpy`` and ``scipy.linalg``
+instances from there without explicitly importing JAX-based libraries.
+
+Enabling Experimental Features
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+To enable experimental features, switch to the ``dev/jax`` branch and use
+
+.. code:: python
+
+   from oqupy.backends import enable_jax_features
+   enable_jax_features()
+
+Alternatively, the `OQUPY_BACKEND` environmental variable may be set to `jax` to
+initialize the jax backend by default.
+
+Contributing Guidelines
+~~~~~~~~~~~~~~~~~~~~~~~
+
+To contribute features compatible with the JAX backend,
+please adhere to the following set of guidelines:
+
+-  avoid wildcard imports of NumPy and SciPy.
+-  use ``from oqupy.backends.numerical_backend import np`` instead of
+   ``import numpy as np`` and use the alias ``default_np`` in cases
+   vanilla NumPy is explicitly required.
+-  use ``from oqupy.backends.numerical_backend import la`` instead of
+   ``import scipy.linalg as la``, except that for non-symmetric
+   eigen-decomposition, ``scipy.linalg.eig`` should be used.
+-  use one of ``np.dtype_complex`` (``np.dtype_float``) or
+   ``oqupy.config.NumPyDtypeComplex`` (``oqupy.config.NumPyDtypeFloat``)
+   instead of ``np.complex_`` (``np.float_``).
+-  convert lists or tuples to arrays when passing them as arguments
+   inside functions.
+-  use ``array = np.update(array, indices, values)`` instead of
+   ``array[indices] = values``.
+-  use ``np.get_random_floats(seed, shape)`` instead of
+   ``np.random.default_rng(seed).random(shape)``.
+-  declare signatures for ``np.vectorize`` explicitly.
+-  avoid directly changing the ``shape`` attribute of an array (use
+   ``.reshape`` instead)
diff --git a/examples/simple_dynamics_jax.py b/examples/simple_dynamics_jax.py
new file mode 100644
index 0000000..5a7f8d1
--- /dev/null
+++ b/examples/simple_dynamics_jax.py
@@ -0,0 +1,43 @@
+#!/usr/bin/env python
+
+import sys
+sys.path.insert(0, '.')
+# set the 'OQUPY_BACKEND' environment variable
+# to 'jax' to initialize JAX backend by default
+# or switch to JAX backend using oqupy.backends
+import oqupy
+from oqupy.backends import enable_jax_features
+# import NumPy from numerical_backend
+#from oqupy.backends.numerical_backend import np
+#enable_jax_features()
+
+import matplotlib.pyplot as plt
+sigma_x = oqupy.operators.sigma("x")
+sigma_z = oqupy.operators.sigma("z")
+up_density_matrix = oqupy.operators.spin_dm("z+")
+Omega = 1.0
+omega_cutoff = 5.0
+alpha = 0.3
+
+system = oqupy.System(0.5 * Omega * sigma_x)
+correlations = oqupy.PowerLawSD(alpha=alpha,
+                                zeta=1,
+                                cutoff=omega_cutoff,
+                                cutoff_type='exponential')
+bath = oqupy.Bath(0.5 * sigma_z, correlations)
+tempo_parameters = oqupy.TempoParameters(dt=0.1, tcut=3.0, epsrel=10**(-4))
+
+dynamics = oqupy.tempo_compute(system=system,
+                               bath=bath,
+                               initial_state=up_density_matrix,
+                               start_time=0.0,
+                               end_time=2.0,
+                               parameters=tempo_parameters,
+                               unique=True)
+t, s_z = dynamics.expectations(0.5*sigma_z, real=True)
+print(s_z)
+plt.plot(t, s_z, label=r'$\alpha=0.3$')
+plt.xlabel(r'$t\,\Omega$')
+plt.ylabel(r'$\langle\sigma_z\rangle$')
+plt.savefig('simple_dynamics_jax.png')
+#plt.show()
diff --git a/oqupy/backends/__init__.py b/oqupy/backends/__init__.py
index e69de29..393f09a 100644
--- a/oqupy/backends/__init__.py
+++ b/oqupy/backends/__init__.py
@@ -0,0 +1,9 @@
+"""Module to initialize OQuPy's backends."""
+
+from oqupy.backends.numerical_backend import set_numerical_backends
+
+def enable_jax_features():
+    """Function to enable experimental features."""
+
+    # set numerical backend to JAX
+    set_numerical_backends('jax')
diff --git a/oqupy/backends/node_array.py b/oqupy/backends/node_array.py
index dc2639a..02f6681 100644
--- a/oqupy/backends/node_array.py
+++ b/oqupy/backends/node_array.py
@@ -15,11 +15,11 @@
 
 from typing import Any, List, Optional, Text, Tuple, Union
 
-import numpy as np
 import tensornetwork as tn
 from tensornetwork import Node
 from tensornetwork.backends.base_backend import BaseBackend
 
+from oqupy.backends.numerical_backend import np
 
 class NodeArray:
     """NodeArray class. """
diff --git a/oqupy/backends/numerical_backend.py b/oqupy/backends/numerical_backend.py
new file mode 100644
index 0000000..39eb88a
--- /dev/null
+++ b/oqupy/backends/numerical_backend.py
@@ -0,0 +1,166 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Module containing NumPy-like and SciPy-like numerical backends.
+"""
+
+import os
+
+import numpy as default_np
+import scipy.linalg as default_la
+
+from tensornetwork.backend_contextmanager import \
+    set_default_backend
+
+import oqupy.config as oc
+
+# store instances of the initialized backends
+# this way, `oqupy.config` remains unchanged
+# and `ocupy.config.DEFAULT_BACKEND` is used
+# when NumPy and LinAlg are initialized
+NUMERICAL_BACKEND_INSTANCES = {}
+
+def get_numerical_backends(
+        backend_name: str,
+    ):
+    """Function to get numerical backend.
+
+    Parameters
+    ----------
+    backend_name: str
+        Name of the backend. Options are `'jax'` and `'numpy'`.
+
+    Returns
+    -------
+    backends: list
+        NumPy and LinAlg backends.
+    """
+
+    _bn = backend_name.lower()
+    if _bn in NUMERICAL_BACKEND_INSTANCES:
+        set_default_backend(_bn)
+        return NUMERICAL_BACKEND_INSTANCES[_bn]
+    assert _bn in ['jax', 'numpy'], \
+        "currently supported backends are `'jax'` and `'numpy'`"
+
+    if 'jax' in _bn:
+        try:
+            # explicitly import and configure jax
+            import jax
+            import jax.numpy as jnp
+            import jax.scipy.linalg as jla
+            jax.config.update('jax_enable_x64', True)
+
+            # # TODO: GPU memory allocation (default is 0.75)
+            # os.environ['XLA_PYTHON_CLIENT_ALLOCATOR'] = 'platform'
+            # os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '.5'
+
+            # set TensorNetwork backend
+            set_default_backend('jax')
+
+            NUMERICAL_BACKEND_INSTANCES['jax'] = [jnp, jla]
+            return NUMERICAL_BACKEND_INSTANCES['jax']
+        except ImportError:
+            print("JAX not installed, defaulting to NumPy")
+
+    # set TensorNetwork backend
+    set_default_backend('numpy')
+
+    NUMERICAL_BACKEND_INSTANCES['numpy'] = [default_np, default_la]
+    return NUMERICAL_BACKEND_INSTANCES['numpy']
+
+class NumPy:
+    """
+    The NumPy backend employing
+    dynamic switching through `oqupy.config`.
+    """
+    def __init__(self,
+                 backend_name=oc.DEFAULT_BACKEND,
+                 ):
+        """Getter for the backend."""
+        self.backend = get_numerical_backends(backend_name)[0]
+
+    @property
+    def dtype_complex(self) -> default_np.dtype:
+        """Getter for the complex datatype."""
+        return oc.NumPyDtypeComplex
+
+    @property
+    def dtype_float(self) -> default_np.dtype:
+        """Getter for the float datatype."""
+        return oc.NumPyDtypeFloat
+
+    def __getattr__(self,
+                    name: str,
+                    ):
+        """Return the backend's default attribute."""
+        return getattr(self.backend, name)
+
+    def update(self,
+               array,
+               indices: tuple,
+               values,
+               ) -> default_np.ndarray:
+        """Option to update select indices of an array with given values."""
+        if not isinstance(array, default_np.ndarray):
+            return array.at[indices].set(values)
+        array[indices] = values
+        return array
+
+    def get_random_floats(self,
+                          seed,
+                          shape,
+                          ):
+        """Method to obtain random floats with a given seed and shape."""
+        random_floats = default_np.random.default_rng(seed).random(shape, \
+                        dtype=default_np.float64)
+        return self.backend.array(random_floats, dtype=self.dtype_float)
+
+class LinAlg:
+    """
+    The Linear Algebra backend employing
+    dynamic switching through `oqupy.config`.
+    """
+    def __init__(self,
+                 backend_name=oc.DEFAULT_BACKEND,
+                 ):
+        """Getter for the backend."""
+        self.backend = get_numerical_backends(backend_name)[1]
+
+    def __getattr__(self,
+                    name: str,
+                    ):
+        """Return the backend's default attribute."""
+        return getattr(self.backend, name)
+
+# setup libraries using environment variable
+# fall back to oqupy.config.DEFAULT_BACKEND
+try:
+    BACKEND_NAME = os.environ[oc.BACKEND_ENV_VAR]
+except KeyError:
+    BACKEND_NAME = oc.DEFAULT_BACKEND
+np = NumPy(backend_name=BACKEND_NAME)
+la = LinAlg(backend_name=BACKEND_NAME)
+
+def set_numerical_backends(
+        backend_name: str
+    ):
+    """Function to set numerical backend.
+
+    Parameters
+    ----------
+    backend_name: str
+        Name of the backend. Options are `'jax'` and `'numpy'`.
+    """
+    backends = get_numerical_backends(backend_name)
+    np.backend = backends[0]
+    la.backend = backends[1]
diff --git a/oqupy/backends/pt_tebd_backend.py b/oqupy/backends/pt_tebd_backend.py
index a018aa6..e5c1b15 100644
--- a/oqupy/backends/pt_tebd_backend.py
+++ b/oqupy/backends/pt_tebd_backend.py
@@ -23,13 +23,13 @@
 import concurrent
 from typing import Dict, List, Tuple
 
-import numpy as np
 from numpy import ndarray
 import tensornetwork as tn
 
-from oqupy.config import NpDtype
+from oqupy.mps_mpo import GateLayer, NnGate, SiteGate
 from oqupy.process_tensor import BaseProcessTensor
-from oqupy.mps_mpo import GateLayer, SiteGate, NnGate
+
+from oqupy.backends.numerical_backend import np
 
 NoneType = type(None)
 
@@ -96,10 +96,12 @@ def __init__(
 
         left_dim = gammas[0].shape[0]
         right_dim = gammas[-1].shape[3]
-        self._lambdas.append(tn.Node(np.identity(left_dim, dtype=NpDtype)))
+        self._lambdas.append(tn.Node(np.identity(left_dim, \
+                                                 dtype=np.dtype_complex)))
         for lam in lambdas:
             self._lambdas.append(tn.Node(np.diag(lam)))
-        self._lambdas.append(tn.Node(np.identity(right_dim, dtype=NpDtype)))
+        self._lambdas.append(tn.Node(np.identity(right_dim, \
+                                                 dtype=np.dtype_complex)))
 
         self._left_e = self._lambdas[0][0]
         self._right_e = self._lambdas[-1][1]
@@ -310,7 +312,7 @@ def _compute_full_trace_gammas(self):
             phys_edge_dim = self._bath_trace_gammas[site].get_dimension(1)
             hs_dim = _isqrt(phys_edge_dim)
             trace_cap = tn.Node(
-                np.identity(hs_dim, dtype=NpDtype).reshape(hs_dim**2))
+                np.identity(hs_dim, dtype=np.dtype_complex).reshape(hs_dim**2))
 
             phys_edge ^ trace_cap[0]
             full_trace_gamma = node_dict[self._bath_trace_gammas[site]] \
@@ -325,7 +327,7 @@ def _compute_total_traces(self):
         nodes = self._full_trace_gammas + self._lambdas
         node_dict, edge_dict = tn.copy(nodes)
 
-        temp_left = tn.Node(np.array([1.0], dtype=NpDtype))
+        temp_left = tn.Node(np.array([1.0], dtype=np.dtype_complex))
         temp_left[0] ^ edge_dict[self._left_e]
         total_left_trace =  temp_left @ node_dict[self._lambdas[0]]
         self._total_left_traces = [total_left_trace.copy()]
@@ -336,7 +338,7 @@ def _compute_total_traces(self):
             tr_gam[1] ^ lam[0]
             total_left_trace = total_left_trace @ tr_gam @ lam
             if site == self._n - 1:
-                temp_right = tn.Node(np.array([1.0], dtype=NpDtype))
+                temp_right = tn.Node(np.array([1.0], dtype=np.dtype_complex))
                 total_left_trace[0] ^ temp_right[0]
                 self._total_trace = (total_left_trace @ temp_right).get_tensor()
             else:
@@ -345,7 +347,7 @@ def _compute_total_traces(self):
         nodes = self._full_trace_gammas + self._lambdas
         node_dict, edge_dict = tn.copy(nodes)
 
-        temp_right = tn.Node(np.array([1.0], dtype=NpDtype))
+        temp_right = tn.Node(np.array([1.0], dtype=np.dtype_complex))
         edge_dict[self._right_e] ^ temp_right[0]
         total_right_trace =  node_dict[self._lambdas[-1]] @ temp_right
         self._total_right_traces = [total_right_trace.copy()]
diff --git a/oqupy/backends/pt_tempo_backend.py b/oqupy/backends/pt_tempo_backend.py
index a2086ad..bedbe5c 100644
--- a/oqupy/backends/pt_tempo_backend.py
+++ b/oqupy/backends/pt_tempo_backend.py
@@ -15,15 +15,13 @@
 
 from typing import Callable, Dict, Optional, List
 
-import numpy as np
-from numpy import ndarray, zeros
-from numpy import max as numpy_max
+from numpy import ndarray
 
 from oqupy.backends import node_array as na
-from oqupy.config import NpDtype
 from oqupy.process_tensor import BaseProcessTensor
-from oqupy import util
+from oqupy.util import add_singleton, create_delta
 
+from oqupy.backends.numerical_backend import np
 
 class PtTempoBackend:
     """
@@ -113,48 +111,56 @@ def initialize(self) -> None:
 
         if self._degeneracy_maps is not None:
             north_degeneracy_map, west_degeneracy_map = self._degeneracy_maps
-            tmp_north_deg_num_vals = numpy_max(north_degeneracy_map)+1
-            tmp_west_deg_num_vals = numpy_max(west_degeneracy_map)+1
+            tmp_north_deg_num_vals = np.max(north_degeneracy_map) + 1
+            tmp_west_deg_num_vals = np.max(west_degeneracy_map) + 1
 
         influences_mpo = []
         influences_mps = []
-        for i in range(self._num_infl):
-            if i == 0:
-                infl = self._influence(i)
-                infl = infl / scale
-                if self._degeneracy_maps is not None:
-                    tmp_mpo = zeros((tmp_west_deg_num_vals,
-                                     self._dimension**2,
-                                     tmp_north_deg_num_vals),
-                                    dtype=complex)
-                    tmp_mps = zeros((self._dimension**2,
-                                     tmp_north_deg_num_vals),
-                                    dtype=complex)
-                    for i1 in range(self._dimension**2):
-                        tmp_mpo[west_degeneracy_map[i1]][i1]\
-                            [north_degeneracy_map[i1]] = \
-                            infl[north_degeneracy_map[i1]]
-                        tmp_mps[i1][north_degeneracy_map[i1]] = \
-                            infl[north_degeneracy_map[i1]]/ scale
-                    infl_mpo = tmp_mpo
-                    infl_mps = tmp_mps
-                else:
-                    infl_mpo = util.create_delta(infl, [1, 1, 0])
-                    infl_mps = infl.T / scale
-            elif i == self._num_infl-1:
-                infl = self._influence(i)
-                infl_mpo = util.add_singleton(infl, 1)
-                infl_mpo = util.add_singleton(infl_mpo, 3)
-                infl_mps = util.add_singleton(infl, 2)
-            else:
-                infl = self._influence(i)
-                infl_mpo = util.create_delta(infl, [0, 1, 1, 0])
-                infl_mps = util.create_delta(infl / scale, [0, 1, 0])
-
-            influences_mpo.append(infl_mpo)
-            influences_mps.append(infl_mps)
-
 
+        # this block takes care of `i == 0`
+        infl = self._influence(0)
+        infl = infl / scale
+        if self._degeneracy_maps is not None:
+            infl_mpo = np.zeros((tmp_west_deg_num_vals, self._dimension**2,
+                                 tmp_north_deg_num_vals), \
+                                    dtype=np.dtype_complex)
+            infl_mps = np.zeros((self._dimension**2,
+                                 tmp_north_deg_num_vals), \
+                                    dtype=np.dtype_complex)
+            # a little bit of optimization is done here by
+            # removing the `for` loop and updating slices
+            _idxs = np.array(list(range(self._dimension**2)))
+            indices = (west_degeneracy_map[_idxs], _idxs,
+                       north_degeneracy_map[_idxs])
+            influences_mpo.append(np.update(
+                array=infl_mpo,
+                indices=indices,
+                values=infl[indices[2]]
+            ))
+            influences_mps.append(np.update(
+                array=infl_mps,
+                indices=indices[1:],
+                values=infl[indices[2]] / scale
+            ))
+        else:
+            influences_mpo.append(create_delta(infl, [1, 1, 0]))
+            influences_mps.append(infl.T / scale)
+
+        # this block takes care of `i > 0` and `i < self._num_infl - 1`
+        # the inner `for` loop can be optimized for parallelization
+        if self._num_infl > 2:
+            indices = list(range(1, self._num_infl - 1))
+            for index in indices:
+                infl = self._influence(index)
+                influences_mpo.append(create_delta(infl, [0, 1, 1, 0]))
+                influences_mps.append(create_delta(infl / scale, [0, 1, 0]))
+
+        # this block takes care of `i == self._num_infl - 1`
+        if self._num_infl > 1:
+            infl = self._influence(self._num_infl - 1)
+            infl_mpo = add_singleton(infl, 1)
+            influences_mpo.append(add_singleton(infl_mpo, 3))
+            influences_mps.append(add_singleton(infl, 2))
 
         self._mpo = na.NodeArray(influences_mpo,
                                  left=False,
@@ -180,7 +186,7 @@ def initialize(self) -> None:
                             max_truncation_err=self._epsrel,
                             relative=True)
 
-        one = np.array([[1.0]], dtype=NpDtype)
+        one = np.array([[1.0]], dtype=np.dtype_complex)
         self._one_na = na.NodeArray([one],
                                    left=True,
                                    right=False,
@@ -239,8 +245,8 @@ def compute_step(self) -> None:
                 dk = int(0 - self._step)
                 infl = self._influence(dk)
                 if infl is not None:
-                    infl_mpo = util.add_singleton(infl, 1)
-                    infl_mpo = util.add_singleton(infl_mpo, 3)
+                    infl_mpo = add_singleton(infl, 1)
+                    infl_mpo = add_singleton(infl_mpo, 3)
                     last_mpo = na.NodeArray(
                             [infl_mpo],
                             left=True,
@@ -282,7 +288,7 @@ def compute_step(self) -> None:
         return self._step < self._num_steps
 
     def get_mpo_tensor(self, step: int) -> ndarray:
-        """ToDo. """
+        """TODO. """
         n = len(self._mps.nodes)
         assert n == self._num_steps
         assert step < n
@@ -290,12 +296,12 @@ def get_mpo_tensor(self, step: int) -> ndarray:
         if step == 0:
             order = [self._mps.bond_edges[0],self._mps.array_edges[0][0]]
             first_t = self._mps.nodes[0].reorder_edges(order).get_tensor()
-            first_t = util.add_singleton(first_t, 0)
+            first_t = add_singleton(first_t, 0)
             tensor = first_t * self._dimension
         elif step == n-1:
             order = [self._mps.bond_edges[-1],self._mps.array_edges[-1][0]]
             last_t = self._mps.nodes[-1].reorder_edges(order).get_tensor()
-            last_t = util.add_singleton(last_t, 1)
+            last_t = add_singleton(last_t, 1)
             tensor = last_t * self._dimension
         else:
             order = [self._mps.bond_edges[step-1],
diff --git a/oqupy/backends/tempo_backend.py b/oqupy/backends/tempo_backend.py
index e5e3015..65b1395 100644
--- a/oqupy/backends/tempo_backend.py
+++ b/oqupy/backends/tempo_backend.py
@@ -12,25 +12,24 @@
 """
 Module for tempo and mean-field tempo backend.
 """
-import os
+
 from copy import copy, deepcopy
 from functools import lru_cache
+import os
 from typing import Callable, Dict, List, Optional, Tuple
 
-import numpy as np
-from numpy import ndarray, moveaxis, expand_dims, eye, exp, outer,\
-    diag, array, kron, swapaxes, reshape, amax, argmax, dot, zeros
-from numpy import max as numpy_max
-from scipy.linalg import svd, LinAlgError
+from numpy import ndarray
 
-from oqupy import operators as op
-from oqupy.config import TEMPO_BACKEND_CONFIG
 from oqupy.backends import node_array as na
+from oqupy.config import TEMPO_BACKEND_CONFIG
+from oqupy import operators as op
 from oqupy.util import create_delta
 
+from oqupy.backends.numerical_backend import np, la
+
 os.environ['NUMPY_EXPERIMENTAL_ARRAY_FUNCTION'] = '0'
 
-class TIBaseBackend: # ToDo: translate this numpy code into tensornetwork
+class TIBaseBackend: # TODO: translate this numpy code into tensornetwork
     """
     Backend class for TEMPO.
 
@@ -74,14 +73,14 @@ def __init__(
         self._max_step = max_step if max_step is None else max_step
         self._kmax = max_step if max_mps_length is None else max_mps_length
         self._initial_data = \
-            eye(self._dim) if initial_data is None else initial_data
+            np.eye(self._dim) if initial_data is None else initial_data
         self._config = TEMPO_BACKEND_CONFIG if config is None else config
 
         self._h_ind, self._h_proj = self._unique(self._ops[0])
         self._v_ind, self._v_proj = self._unique(zip(*self._ops[1:]))
         self._v_dim, self._h_dim, = len(self._v_ind), len(self._h_ind)
 
-        self._cap = expand_dims(eye(self._dim), -1)
+        self._cap = np.expand_dims(np.eye(self._dim), -1)
         self._step = None
         self._mps = None
         self.data = [self._initial_data]
@@ -127,18 +126,19 @@ def _influence_tensor(self, k):
         if k == 0:
             c0 = self._coefficients(0)
             o0_2 = c0.real * self._ops[1] - 1j * c0.imag * self._ops[2]
-            tensor = exp(outer(o_2, o_1)) * np.dot(prop, prop).T \
-                * exp(o_1 * o0_2)
+            tensor = np.exp(np.outer(o_2, o_1)) * np.dot(prop, prop).T \
+                * np.exp(o_1 * o0_2)
             tensor = np.dot(
-                kron(self._v_proj, self._h_proj), diag(tensor.flatten()))
-            tensor = reshape(
+                np.kron(self._v_proj, self._h_proj), np.diag(tensor.flatten()))
+            tensor = np.reshape(
                 tensor, (self._v_dim, self._h_dim, self._dim, self._dim))
-            tensor = moveaxis(swapaxes(tensor, 2, 3), 0, 2)
+            tensor = np.moveaxis(np.swapaxes(tensor, 2, 3), 0, 2)
         else:
-            tensor = diag(exp(kron(o_2[self._v_ind], o_1[self._h_ind])))
-            tensor = reshape(
+            tensor = np.diag(np.exp(np.kron(o_2[self._v_ind], \
+                                            o_1[self._h_ind])))
+            tensor = np.reshape(
                 tensor, (self._v_dim, self._h_dim, self._v_dim, self._h_dim))
-            tensor = swapaxes(tensor, 0, 3)
+            tensor = np.swapaxes(tensor, 0, 3)
         return tensor
 
 
@@ -158,7 +158,7 @@ def _contract(self, k) -> List:
         tens = self._influence_tensor(len(self._mps) - 1 - k)
         if k == 0:
             # sum over mpo_w and create new leg with mpo_w=1
-            tens = expand_dims(tens.sum(0), 0)
+            tens = np.expand_dims(tens.sum(0), 0)
 
         mps_w, _, mps_e = self._mps[k].shape
         mpo_w, mpo_e, mpo_n, _ = tens.shape
@@ -176,10 +176,10 @@ def _truncate_right(self, k) -> int:
         site."""
         west, north, east = self._mps[k].shape
         u, s, v_dag = self._scipy_svd(
-            reshape(self._mps[k], (west * north, east)), self._precision)
-        self._mps[k] = reshape(u, (west, north, len(s)))
+            np.reshape(self._mps[k], (west * north, east)), self._precision)
+        self._mps[k] = np.reshape(u, (west, north, len(s)))
         self._mps[k + 1] = np.dot(
-            (s * v_dag.T).T, swapaxes(self._mps[k + 1], 0, 1))
+            (s * v_dag.T).T, np.swapaxes(self._mps[k + 1], 0, 1))
         return len(s)
 
     def _truncate_left(self, k) -> int:
@@ -188,14 +188,14 @@ def _truncate_left(self, k) -> int:
         self._mps[k] = self._mps[k].T
         east, north, west = self._mps[k].shape
         u, s, v_dag = self._scipy_svd(
-            reshape(self._mps[k], (east * north, west)), self._precision)
-        self._mps[k] = reshape(u, (east, north, len(s))).T
+            np.reshape(self._mps[k], (east * north, west)), self._precision)
+        self._mps[k] = np.reshape(u, (east, north, len(s))).T
         self._mps[k - 1] = np.dot(self._mps[k - 1], v_dag.T * s)
         return len(s)
 
 
     def initialise(self, step=None, mps=None) -> Tuple[int, ndarray]:
-        """ ToDo """
+        """ TODO """
         if mps is not None:
             self._mps = mps
             self._step = step
@@ -205,17 +205,18 @@ def initialise(self, step=None, mps=None) -> Tuple[int, ndarray]:
             o_1 = self._ops[0]
             o_2 = c_real * self._ops[1] - 1j * c_imag * self._ops[2]
             # shape ( e, n, s)
-            tensor = np.dot(self._initial_data, self._prop.T * exp(o_1 * o_2))
+            tensor = np.dot(self._initial_data, self._prop.T * \
+                            np.exp(o_1 * o_2))
             self.data.append(np.dot(tensor, self._prop.T))
             # contains whole timestep freeprop!
             tensor = np.dot(self._influence_tensor(0), tensor.T)
-            tensor = swapaxes(tensor.sum(0), 0, 2)
+            tensor = np.swapaxes(tensor.sum(0), 0, 2)
             self._mps = [tensor, self._cap]
             self.data.append(self.readout())
             self._step = 1
         return self._step, self.data[-1]
 
-    def compute_step(self) -> Tuple[int, ndarray]: # ToDo: make readout optional
+    def compute_step(self) -> Tuple[int, ndarray]: # TODO: make readout optional
         """
         Takes a step in the TEMPO tensor network computation.
 
@@ -263,11 +264,11 @@ def compute_step(self) -> Tuple[int, ndarray]: # ToDo: make readout optional
         # turn east leg at last site into north leg at new last site
         self._mps.append(self._cap)
 
-        if len(self._mps) > self._kmax + 1:  # ToDo: check this
+        if len(self._mps) > self._kmax + 1:  # TODO: check this
             # remove first site, turn into matrix
             end = self._mps.pop(0).sum(1)
             # np.dot into new first site
-            self._mps[0] = np.dot(end, swapaxes(self._mps[0], 0, 1))
+            self._mps[0] = np.dot(end, np.swapaxes(self._mps[0], 0, 1))
 
         self._step += 1
         self.data.append(self.readout())
@@ -293,12 +294,12 @@ def _scipy_svd(theta, precision):
         """ static svd truncation method using scipy.linalg.svd"""
         try:
             u, singular_values, v_dagger = \
-                svd(theta, full_matrices=False, lapack_driver='gesvd')
-        except LinAlgError:
+                la.svd(theta, full_matrices=False, lapack_driver='gesvd')
+        except la.LinAlgError:
             print('svd except')
             u, singular_values, v_dagger = \
-                svd(theta, full_matrices=False, lapack_driver='gesdd')
-        chi = argmax(singular_values/amax(singular_values) < precision)
+                la.svd(theta, full_matrices=False, lapack_driver='gesdd')
+        chi = np.argmax(singular_values/np.amax(singular_values) < precision)
         if not chi:
             chi = len(singular_values)
         return u[:, 0:chi], singular_values[0:chi], v_dagger[0:chi, :]
@@ -307,8 +308,8 @@ def _scipy_svd(theta, precision):
     def _unique(values):
         vals = list(values)
         inverse = [vals.index(e) for e in vals]
-        indices = array(sorted(set(inverse), key=inverse.index))
-        inverse = array([[int(i == j) for i in inverse] for j in indices])
+        indices = np.array(sorted(set(inverse), key=inverse.index))
+        inverse = np.array([[int(i == j) for i in inverse] for j in indices])
         return indices, inverse
 
 class BaseTempoBackend:
@@ -377,7 +378,7 @@ def step(self) -> int:
         return self._step
 
     def initialize_mps_mpo(self):
-        """ToDo"""
+        """TODO"""
         self._initial_state = copy(self._initial_state).reshape(-1)
 
         self._super_u = op.left_right_super(
@@ -398,43 +399,53 @@ def initialize_mps_mpo(self):
             dkmax_pre_compute = self._dkmax + 1
 
         if self._degeneracy_maps is not None:
-            north_degeneracy_map, west_degeneracy_map =\
+            north_degeneracy_map, west_degeneracy_map = \
                     self._degeneracy_maps
-            tmp_north_deg_num_vals = numpy_max(north_degeneracy_map)+1
-            tmp_west_deg_num_vals = numpy_max(west_degeneracy_map)+1
-
-        for i in range(dkmax_pre_compute):
-            infl = self._influence(i)
-            if i == 0:
-                if self._degeneracy_maps is not None:
-                    tmp=zeros((tmp_west_deg_num_vals,self._dim**2,
-                               tmp_north_deg_num_vals,self._dim**2),
-                              dtype=complex)
-                    for i1 in range(self._dim**2):
-                        tmp[west_degeneracy_map[i1]][i1] \
-                            [north_degeneracy_map[i1]][i1]= \
-                                infl[north_degeneracy_map[i1]]
-                    infl_four_legs = tmp
-                else:
-                    infl_four_legs = create_delta(infl, [1, 0, 0, 1])
-                tmp = dot(moveaxis(infl_four_legs, 1, -1),
-                        self._super_u_dagg)
-                tmp = moveaxis(tmp, -1, 1)
-                tmp = np.dot(tmp, self._super_u.T)
-                infl_four_legs = tmp
-            else:
-                infl_four_legs = create_delta(infl, [1, 0, 0, 1])
+            tmp_north_deg_num_vals = np.max(north_degeneracy_map) + 1
+            tmp_west_deg_num_vals = np.max(west_degeneracy_map) + 1
 
-            influences.append(infl_four_legs)
+        influences = []
+
+        # this block takes care of `i == 0`
+        infl = self._influence(0)
+        if self._degeneracy_maps is not None:
+            infl_four_legs = np.zeros((tmp_west_deg_num_vals, self._dim**2,
+                                       tmp_north_deg_num_vals, self._dim**2), \
+                                        dtype=np.dtype_complex)
+            # a little bit of optimization is done here by
+            # removing the `for` loop and updating slices
+            _idxs = np.array(list(range(self._dim**2)))
+            indices = (west_degeneracy_map[_idxs], _idxs,
+                       north_degeneracy_map[_idxs], _idxs)
+            infl_four_legs = np.update(
+                array=infl_four_legs,
+                indices=indices,
+                values=infl[indices[2]]
+            )
+        else:
+            infl_four_legs = create_delta(infl, [1, 0, 0, 1])
+        infl_four_legs = np.dot(np.moveaxis(infl_four_legs, 1, -1), \
+                                self._super_u_dagg)
+        infl_four_legs = np.moveaxis(infl_four_legs, -1, 1)
+        infl_four_legs = np.dot(infl_four_legs, self._super_u.T)
+        influences.append(infl_four_legs)
+
+        # this block takes care of `i > 0`
+        # the inner `for` loop can be optimized for parallelization
+        if dkmax_pre_compute > 1:
+            indices = list(range(1, dkmax_pre_compute))
+            for index in indices:
+                infl = self._influence(index)
+                influences.append(create_delta(infl, [1, 0, 0, 1]))
 
-        self._mps = na.NodeArray([self._initial_state],
-                                 left=False,
-                                 right=False,
-                                 name="Thee MPS")
         self._mpo = na.NodeArray(list(reversed(influences)),
                                  left=True,
                                  right=True,
                                  name="Thee Time Evolving MPO")
+        self._mps = na.NodeArray([self._initial_state],
+                                 left=False,
+                                 right=False,
+                                 name="Thee MPS")
 
     def compute_system_step(self, current_step, prop_1, prop_2) -> ndarray:
         """
@@ -577,7 +588,7 @@ def compute_system_step(self, current_step, prop_1, prop_2) -> ndarray:
 
 class TempoBackend(BaseTempoBackend):
     """
-    ToDo
+    TODO
     """
 
     def __init__(
@@ -609,7 +620,7 @@ def __init__(
 
     def initialize(self) -> Tuple[int, ndarray]:
         """
-        ToDo
+        TODO
         """
         self._step = 0
         self.initialize_mps_mpo()
@@ -618,7 +629,7 @@ def initialize(self) -> Tuple[int, ndarray]:
 
     def compute_step(self) -> Tuple[int, ndarray]:
         """
-        ToDo
+        TODO
         """
         self._step += 1
         prop_1, prop_2 = self._propagators(self._step - 1)
diff --git a/oqupy/base_api.py b/oqupy/base_api.py
index d89dbbf..b8dc1c0 100644
--- a/oqupy/base_api.py
+++ b/oqupy/base_api.py
@@ -17,7 +17,6 @@
 
 from oqupy.config import SEPERATOR
 
-
 class BaseAPIClass:
     """
     Base class for API objects
diff --git a/oqupy/bath.py b/oqupy/bath.py
index 4266297..16b7517 100644
--- a/oqupy/bath.py
+++ b/oqupy/bath.py
@@ -12,18 +12,18 @@
 """
 Module on physical information on the bath and its coupling to the system.
 """
-from typing import Optional, Text
+
 from copy import copy
+from typing import Optional, Text
 
-import numpy as np
 from numpy import ndarray
 
-from oqupy.config import NpDtype
-from oqupy.config import DEFAULT_TOLERANCE_DEGENERACY
 from oqupy.bath_correlations import BaseCorrelations
 from oqupy.base_api import BaseAPIClass
+from oqupy.config import DEFAULT_TOLERANCE_DEGENERACY
 from oqupy.operators import commutator, acommutator
 
+from oqupy.backends.numerical_backend import np
 
 class Bath(BaseAPIClass):
     """
@@ -50,8 +50,11 @@ def __init__(
         """Creates a Bath object. """
         # input check for coupling_operator.
         try:
-            tmp_coupling_operator = np.array(coupling_operator, dtype=NpDtype)
-            tmp_coupling_operator.setflags(write=False)
+            tmp_coupling_operator = np.array(coupling_operator, \
+                                             dtype=np.dtype_complex)
+            # set immutable if default NumPy array
+            if isinstance(tmp_coupling_operator, ndarray):
+                tmp_coupling_operator.setflags(write=False)
         except Exception as e:
             raise AssertionError("Coupling operator must be numpy array") \
                 from e
@@ -105,17 +108,17 @@ def __str__(self) -> Text:
         return "".join(ret)
 
     @property
-    def coupling_operator(self) -> np.ndarray:
+    def coupling_operator(self) -> ndarray:
         """The diagonalised system operator to which the bath couples. """
         return self._coupling_operator.copy()
 
     @property
-    def unitary_transform(self) -> np.ndarray:
+    def unitary_transform(self) -> ndarray:
         """The unitary that makes the coupling operator diagonal. """
         return self._unitary.copy()
 
     @property
-    def dimension(self) -> np.ndarray:
+    def dimension(self) -> ndarray:
         """Hilbert space dimension of the coupling operator. """
         return copy(self._dimension)
 
@@ -125,19 +128,19 @@ def correlations(self) -> BaseCorrelations:
         return copy(self._correlations)
 
     @property
-    def coupling_acomm(self) -> np.ndarray:
+    def coupling_acomm(self) -> ndarray:
         """Diagonal elements of the anti-commutator of the coupling
         operator. """
         return self._coupling_acomm.copy()
 
     @property
-    def coupling_comm(self) -> np.ndarray:
+    def coupling_comm(self) -> ndarray:
         """Diagonal elements of the commutator of the coupling
         operator. """
         return self._coupling_comm.copy()
 
     @property
-    def north_degeneracy_map(self) -> np.ndarray:
+    def north_degeneracy_map(self) -> ndarray:
         """Map to minimal set of indices for influence tensors in
         north-south direction according to simultaneous degeneracies in
         sums & differences of eigenvalues of coupling operator (minimal
@@ -146,7 +149,7 @@ def north_degeneracy_map(self) -> np.ndarray:
         return copy(self._north_degeneracy_map)
 
     @property
-    def west_degeneracy_map(self) -> np.ndarray:
+    def west_degeneracy_map(self) -> ndarray:
         """Map to minimal set of indices for influence tensors in
         west-east direction according to degeneracies in sums of
         eigenvalues of coupling operator (minimal dimension is number
@@ -160,5 +163,5 @@ def _row_degeneracy(matrix):
     indices mapping full space to non-degenerate rows (repeated
     indices indicate row degeneracy in the original matrix)."""
     mat = np.array(matrix).round(decimals=DEFAULT_TOLERANCE_DEGENERACY)
-    return_map = np.unique(mat.T,return_inverse=True,axis=0)[1]
+    return_map = np.unique(mat.T,return_inverse=True,axis=0)[1].flatten()
     return return_map
diff --git a/oqupy/bath_correlations.py b/oqupy/bath_correlations.py
index 3bc40b4..0c53f34 100644
--- a/oqupy/bath_correlations.py
+++ b/oqupy/bath_correlations.py
@@ -13,17 +13,18 @@
 Module for environment correlations.
 """
 
+from functools import lru_cache
 from typing import Callable, Optional, Text
 from typing import Any as ArrayLike
-from functools import lru_cache
 
-import numpy as np
 from scipy import integrate
 
 from oqupy.base_api import BaseAPIClass
 from oqupy.config import INTEGRATE_EPSREL, SUBDIV_LIMIT
 from oqupy.util import check_true
 
+from oqupy.backends.numerical_backend import np
+
 #np.seterr(all='warn')
 # --- spectral density classes ------------------------------------------------
 
diff --git a/oqupy/bath_dynamics.py b/oqupy/bath_dynamics.py
index b9b4cd4..f958cde 100644
--- a/oqupy/bath_dynamics.py
+++ b/oqupy/bath_dynamics.py
@@ -19,16 +19,16 @@
 """
 
 from typing import Optional, Text, Tuple
-import numpy as np
+
 from numpy import ndarray
 
 from oqupy.base_api import BaseAPIClass
-from oqupy.process_tensor import BaseProcessTensor
 from oqupy.bath import Bath
+from oqupy.process_tensor import BaseProcessTensor
 from oqupy.system import BaseSystem
-from oqupy.config import NpDtype
 from oqupy.system_dynamics import compute_correlations
 
+from oqupy.backends.numerical_backend import np
 
 class TwoTimeBathCorrelations(BaseAPIClass):
     """
@@ -75,7 +75,7 @@ def __init__(
         self._initial_state = initial_state
 
         if system_correlations is None:
-            self._system_correlations = np.array([[]], dtype=NpDtype)
+            self._system_correlations = np.array([[]], dtype=np.dtype_complex)
         else:
             self._system_correlations = system_correlations
         self._temp = bath.correlations.temperature
@@ -181,8 +181,8 @@ def occupation(
         last_time = corr_mat_dim * dt
         tlist = np.arange(0, last_time+dt, dt)
         if freq == 0:
-            return tlist, np.ones(len(tlist),
-                                  dtype=NpDtype) * (np.nan + 1.0j*np.nan)
+            return tlist, np.ones(len(tlist), dtype=np.dtype_complex) * \
+                (np.nan + 1.0j*np.nan)
         self.generate_system_correlations(last_time, progress_type)
         _sys_correlations = self._system_correlations[:corr_mat_dim,
                                                       :corr_mat_dim]
@@ -195,7 +195,7 @@ def occupation(
             np.sum(_sys_correlations.real*re_kernel \
                           + 1j*_sys_correlations.imag*im_kernel, axis = 0)
                 ).real * coup
-        bath_occupation = np.append([0], bath_occupation)
+        bath_occupation = np.append(np.array([0]), bath_occupation)
         if not change_only and self._temp > 0:
             bath_occupation += np.exp(-freq/self._temp) \
                 / (1 - np.exp(-freq/self._temp))
@@ -352,8 +352,8 @@ def _calc_kernel(self,
         ker_dim = int(np.round(time_2 / dt))
         # calculate index corresponding to t_1
         switch = int(np.round(time_1 / dt))
-        re_kernel = np.zeros((ker_dim, ker_dim), dtype = NpDtype)
-        im_kernel = np.zeros((ker_dim, ker_dim), dtype = NpDtype)
+        re_kernel = np.zeros((ker_dim, ker_dim), dtype = np.dtype_complex)
+        im_kernel = np.zeros((ker_dim, ker_dim), dtype = np.dtype_complex)
 
         tpp_index, tp_index = np.meshgrid(
             np.arange(ker_dim), np.arange(ker_dim),
@@ -396,54 +396,121 @@ def phase(region, swap_ts = False):
                 ph += np.exp(a * tk*dt + b * tkp*dt) / (a * b)
             return ph
 
-
         if dagg == (0, 1):
-            re_kernel[regions['a']] = phase('a') + phase('a', 1)
-
-            re_kernel[regions['b']] = phase('b')
-
-            im_kernel[regions['a']] = ((2*n_1 + 1) * phase('a') -
-                                       (2*n_2 + 1) * phase('a', 1))
-
-            im_kernel[regions['b']] = (2*n_1 + 1) * phase('b')
-
-            im_kernel[regions['c']] = -2 * (n_1 + 1) * phase('c')
+            re_kernel = np.update(
+                array=re_kernel,
+                indices=regions['a'],
+                values=phase('a') + phase('a', 1)
+            )
+            re_kernel = np.update(
+                array=re_kernel,
+                indices=regions['b'],
+                values=phase('b')
+            )
+
+            im_kernel = np.update(
+                array=im_kernel,
+                indices=regions['a'],
+                values=((2 * n_1 + 1) * phase('a') -
+                        (2 * n_2 + 1) * phase('a', 1))
+            )
+            im_kernel = np.update(
+                array=im_kernel,
+                indices=regions['b'],
+                values=(2 * n_1 + 1) * phase('b')
+            )
+            im_kernel = np.update(
+                array=im_kernel,
+                indices=regions['c'],
+                values=- 2 * (n_1 + 1) * phase('c')
+            )
 
         elif dagg == (1, 0):
-            re_kernel[regions['a']] = phase('a') + phase('a', 1)
-
-            re_kernel[regions['b']] = phase('b')
-
-            im_kernel[regions['a']] = ((2*n_1 + 1) * phase('a') -
-                                       (2*n_2 + 1) * phase('a', 1))
-
-            im_kernel[regions['b']] = (2*n_1 + 1) * phase('b')
-
-            im_kernel[regions['c']] = 2 * n_1 * phase('c')
+            re_kernel = np.update(
+                array=re_kernel,
+                indices=regions['a'],
+                values=phase('a') + phase('a', 1)
+            )
+            re_kernel = np.update(
+                array=re_kernel,
+                indices=regions['b'],
+                values=phase('b')
+            )
+
+            im_kernel = np.update(
+                array=im_kernel,
+                indices=regions['a'],
+                values=((2*n_1 + 1) * phase('a') -
+                        (2*n_2 + 1) * phase('a', 1))
+            )
+            im_kernel = np.update(
+                array=im_kernel,
+                indices=regions['b'],
+                values=(2*n_1 + 1) * phase('b')
+            )
+            im_kernel = np.update(
+                array=im_kernel,
+                indices=regions['c'],
+                values=2 * n_1 * phase('c')
+            )
 
         elif dagg == (1, 1):
-            re_kernel[regions['a']] = -(phase('a') + phase('a', 1))
-
-            re_kernel[regions['b']] = -phase('b')
-
-            im_kernel[regions['a']] = ((2*n_1 + 1) * phase('a') +
-                                       (2*n_2 + 1) * phase('a', 1))
-
-            im_kernel[regions['b']] = (2*n_1 + 1) * phase('b')
-
-            im_kernel[regions['c']] = 2 * (n_1 + 1) * phase('c')
+            re_kernel = np.update(
+                array=re_kernel,
+                indices=regions['a'],
+                values=- (phase('a') + phase('a', 1))
+            )
+            re_kernel = np.update(
+                array=re_kernel,
+                indices=regions['b'],
+                values=- phase('b')
+            )
+
+            im_kernel = np.update(
+                array=im_kernel,
+                indices=regions['a'],
+                values=((2*n_1 + 1) * phase('a') +
+                        (2*n_2 + 1) * phase('a', 1))
+            )
+            im_kernel = np.update(
+                array=im_kernel,
+                indices=regions['b'],
+                values=(2*n_1 + 1) * phase('b')
+            )
+            im_kernel = np.update(
+                array=im_kernel,
+                indices=regions['c'],
+                values=2 * (n_1 + 1) * phase('c')
+            )
 
         elif dagg == (0, 0):
-            re_kernel[regions['a']] = -(phase('a') + phase('a', 1))
-
-            re_kernel[regions['b']] = -phase('b')
-
-            im_kernel[regions['a']] = -((2*n_2 + 1) * phase('a', 1) +
-                                        (2*n_1 + 1) * phase('a'))
-
-            im_kernel[regions['b']] = -(2*n_1 + 1) * phase('b')
-
-            im_kernel[regions['c']] = -2 * n_1 * phase('c')
+            re_kernel = np.update(
+                array=re_kernel,
+                indices=regions['a'],
+                values=- (phase('a') + phase('a', 1))
+            )
+            re_kernel = np.update(
+                array=re_kernel,
+                indices=regions['b'],
+                values=- phase('b')
+            )
+
+            im_kernel = np.update(
+                array=im_kernel,
+                indices=regions['a'],
+                values=- ((2*n_1 + 1) * phase('a') +
+                        (2*n_2 + 1) * phase('a', 1))
+            )
+            im_kernel = np.update(
+                array=im_kernel,
+                indices=regions['b'],
+                values=- (2*n_1 + 1) * phase('b')
+            )
+            im_kernel = np.update(
+                array=im_kernel,
+                indices=regions['c'],
+                values=- 2 * n_1 * phase('c')
+            )
 
         re_kernel = np.triu(re_kernel) #only keep triangular region
         im_kernel = np.triu(im_kernel)
diff --git a/oqupy/config.py b/oqupy/config.py
index 592dfa7..f2a6be8 100644
--- a/oqupy/config.py
+++ b/oqupy/config.py
@@ -13,11 +13,12 @@
 Module to define global configuration for the time_evovling_mpo package.
 """
 
-from numpy import float64, complex128
-
-# Numpy datatype
-NpDtype = complex128
-NpDtypeReal = float64
+# numerical backend
+import numpy as default_np
+BACKEND_ENV_VAR = 'OQUPY_BACKEND'
+DEFAULT_BACKEND = 'numpy'
+NumPyDtypeComplex = default_np.complex128   # earlier NpDtype
+NumPyDtypeFloat = default_np.float64        # earlier NpDtypeReal
 
 # Separator string for __str__ functions
 SEPERATOR = "----------------------------------------------\n"
diff --git a/oqupy/control.py b/oqupy/control.py
index 1cbe315..27ee5dc 100644
--- a/oqupy/control.py
+++ b/oqupy/control.py
@@ -18,14 +18,14 @@
 framework and efficient characterization*, Phys. Rev. A 97, 012127 (2018).
 """
 
-from typing import Callable, List, Optional, Text, Tuple, Union
 from copy import deepcopy
+from typing import Callable, List, Optional, Text, Tuple, Union
 
-import numpy as np
 from numpy import ndarray
 
 from oqupy.base_api import BaseAPIClass
-from oqupy.config import NpDtype
+
+from oqupy.backends.numerical_backend import np
 
 class Control(BaseAPIClass):
     """
@@ -112,7 +112,7 @@ def add_continuous(
             control_fct: Callable[[ndarray, float], ndarray],
             post: Optional[bool] = False) -> None:
         """
-        ToDo
+        TODO
         """
         raise NotImplementedError()
 
@@ -252,7 +252,7 @@ def add_single_site_control(
         assert isinstance(site, int)
         assert site < len(self)
         assert isinstance(step, int)
-        contr = np.array(control,  dtype=NpDtype)
+        contr = np.array(control,  dtype=np.dtype_complex)
         assert contr.shape == (self._hs_dims[site]**2, self._hs_dims[site]**2)
         if not post:
             self._single_site_controls_pre.append({
diff --git a/oqupy/dynamics.py b/oqupy/dynamics.py
index f58c737..f657219 100644
--- a/oqupy/dynamics.py
+++ b/oqupy/dynamics.py
@@ -13,15 +13,15 @@
 Module on the discrete time evolution of a density matrix.
 """
 
-from typing import List, Optional, Text, Tuple
-from copy import copy
 from bisect import bisect
+from copy import copy
+from typing import List, Optional, Text, Tuple
 
-import numpy as np
 from numpy import ndarray
 
 from oqupy.base_api import BaseAPIClass
-from oqupy.config import NpDtype, NpDtypeReal
+
+from oqupy.backends.numerical_backend import np
 
 class BaseDynamics(BaseAPIClass):
     """
@@ -52,9 +52,9 @@ def __str__(self) -> Text:
         ret.append("  length        = {} timesteps \n".format(len(self)))
         if len(self) > 0:
             ret.append("  min time      = {} \n".format(
-                np.min(self._times)))
+                np.min(self.times)))
             ret.append("  max time      = {} \n".format(
-                np.max(self._times)))
+                np.max(self.times)))
         return "".join(ret)
 
     def __len__(self) -> int:
@@ -63,12 +63,12 @@ def __len__(self) -> int:
     @property
     def times(self) -> ndarray:
         """Times of the dynamics. """
-        return np.array(self._times, dtype=NpDtypeReal)
+        return np.array(self._times, dtype=np.dtype_float)
 
     @property
     def states(self) -> ndarray:
         """States of the dynamics. """
-        return np.array(self._states, dtype=NpDtype)
+        return np.array(self._states, dtype=np.dtype_complex)
 
     @property
     def shape(self) -> ndarray:
@@ -106,10 +106,10 @@ def expectations(
         if len(self) == 0:
             return None, None
         if operator is None:
-            tmp_operator = np.identity(self._shape[0], dtype=NpDtype)
+            tmp_operator = np.identity(self._shape[0], dtype=np.dtype_complex)
         else:
             try:
-                tmp_operator = np.array(operator, dtype=NpDtype)
+                tmp_operator = np.array(operator, dtype=np.dtype_complex)
             except Exception as e:
                 raise AssertionError("Argument `operator` must be ndarray.") \
                     from e
@@ -181,7 +181,7 @@ def add(
     @property
     def shape(self):
         """The shape of the states. """
-        return np.copy(self._shape)
+        return copy(self._shape)
 
 class MeanFieldDynamics(BaseAPIClass):
     """
@@ -269,12 +269,12 @@ def __len__(self) -> int:
     @property
     def times(self) -> ndarray:
         """Times of the dynamics. """
-        return np.array(self._times, dtype=NpDtypeReal)
+        return np.array(self._times, dtype=np.dtype_float)
 
     @property
     def fields(self) -> ndarray:
         """Fields of the dynamics. """
-        return np.array(self._fields, dtype=NpDtype)
+        return np.array(self._fields, dtype=np.dtype_complex)
 
     def field_expectations(self) -> Tuple[ndarray, ndarray]:
         r"""
@@ -289,8 +289,8 @@ def field_expectations(self) -> Tuple[ndarray, ndarray]:
         """
         if len(self) == 0:
             return None, None
-        return np.array(self._times, dtype=NpDtypeReal),\
-                np.array(self._fields, dtype=NpDtype)
+        return np.array(self._times, dtype=np.dtype_float), \
+                np.array(self._fields, dtype=np.dtype_complex)
 
     @property
     def system_dynamics(self) -> List[Dynamics]:
@@ -336,7 +336,7 @@ def _parse_time(time) -> float:
 
 def _parse_state(state, previous_shape) -> Tuple[ndarray, Tuple[int]]:
     try:
-        tmp_state = np.array(state, dtype=NpDtype)
+        tmp_state = np.array(state, dtype=np.dtype_complex)
     except Exception as e:
         raise AssertionError("Argument `state` must be ndarray.") from e
     tmp_shape = tmp_state.shape
diff --git a/oqupy/gradient.py b/oqupy/gradient.py
index 9c17743..6be0b81 100644
--- a/oqupy/gradient.py
+++ b/oqupy/gradient.py
@@ -13,20 +13,23 @@
 Frontend for computing the gradient of some objective function w.r.t. some
 control parameters.
 """
+
 from typing import Callable, Dict, List, Optional, Text, Tuple, Union
 
-import numpy as np
 from numpy import ndarray
 import tensornetwork as tn
 
-from oqupy.system_dynamics import _compute_dynamics_input_parse, \
-    _apply_system_superoperator, _apply_derivative_pt_mpos, _get_pt_mpos, \
-    _get_pt_mpos_backprop, _get_caps, _apply_caps, _apply_pt_mpos
 from oqupy.control import Control
 from oqupy.dynamics import Dynamics
 from oqupy.process_tensor import BaseProcessTensor
+from oqupy.system_dynamics import _apply_caps, _apply_derivative_pt_mpos, \
+    _apply_pt_mpos, _apply_system_superoperator, \
+    _compute_dynamics_input_parse, _get_caps, \
+    _get_pt_mpos, _get_pt_mpos_backprop
 from oqupy.system import ParameterizedSystem
-from oqupy.util import get_progress, check_isinstance
+from oqupy.util import check_isinstance, get_progress
+
+from oqupy.backends.numerical_backend import np
 
 def state_gradient(
         system: ParameterizedSystem,
@@ -73,7 +76,7 @@ def state_gradient(
         'gradient' : derivatives of Z with respect to the parameters
         'dynamics' : the dynamics of the system
     """
-    check_isinstance(parameters, ndarray, 'parameters')
+    check_isinstance(parameters, np.ndarray, 'parameters')
 
     num_steps = len(process_tensors[0])
     dt = process_tensors[0].dt
@@ -151,14 +154,20 @@ def combine_derivs(
         prog_bar.update(i)
 
         for j in range(0,num_parameters):
-            total_derivs[2*i][j] = combine_derivs(
-                            adjoint_tensor[i],
-                            first_half_prop_derivs[j].T,
-                            second_half_prop.T)
-            total_derivs[2*i+1][j] = combine_derivs(
-                adjoint_tensor[i],
-                first_half_prop.T,
-                second_half_prop_derivs[j].T)
+            total_derivs = np.update(
+                array=total_derivs,
+                indices=(2 * i, j),
+                values=combine_derivs(adjoint_tensor[i],
+                                      first_half_prop_derivs[j].T,
+                                      second_half_prop.T)
+            )
+            total_derivs = np.update(
+                array=total_derivs,
+                indices=(2 * i + 1, j),
+                values=combine_derivs(adjoint_tensor[i],
+                                      first_half_prop.T,
+                                      second_half_prop_derivs[j].T)
+            )
 
     prog_bar.update(num_steps)
     prog_bar.exit()
@@ -262,7 +271,7 @@ def controls(step: int):
     #    edges 0, 1, .., num_envs-1    are the bond legs of the environments
     #    edge  -1                      is the state leg
     initial_ndarray = initial_state.reshape(hs_dim**2)
-    initial_ndarray.shape = tuple([1]*num_envs+[hs_dim**2])
+    initial_ndarray = initial_ndarray.reshape(tuple([1]*num_envs+[hs_dim**2]))
     current_node = tn.Node(initial_ndarray)
     current_edges = current_node[:]
 
@@ -353,7 +362,7 @@ def controls(step: int):
 
     target_ndarray = target_derivative
     target_ndarray = target_ndarray.reshape(hs_dim**2)
-    target_ndarray.shape = tuple([1]*num_envs+[hs_dim**2])
+    target_ndarray = target_ndarray.reshape(tuple([1]*num_envs+[hs_dim**2]))
     current_node = tn.Node(target_ndarray)
     current_edges = current_node[:]
 
diff --git a/oqupy/helpers.py b/oqupy/helpers.py
index 7127c1d..9969a8b 100644
--- a/oqupy/helpers.py
+++ b/oqupy/helpers.py
@@ -13,13 +13,13 @@
 Handy helper functions.
 """
 
-import numpy as np
-import matplotlib.pyplot as plt
 from matplotlib.axes import Axes
+import matplotlib.pyplot as plt
 
 from oqupy.bath_correlations import BaseCorrelations
 from oqupy.tempo import TempoParameters
 
+from oqupy.backends.numerical_backend import np
 
 def plot_correlations_with_parameters(
         correlations: BaseCorrelations,
diff --git a/oqupy/mps_mpo.py b/oqupy/mps_mpo.py
index 580a158..a18b760 100644
--- a/oqupy/mps_mpo.py
+++ b/oqupy/mps_mpo.py
@@ -16,12 +16,12 @@
 from copy import deepcopy
 from typing import Any, List, Optional, Text, Tuple
 
-import numpy as np
 from numpy import ndarray
-from scipy import linalg
 
 from oqupy.base_api import BaseAPIClass
-from oqupy.config import NpDtype
+
+from oqupy.backends.numerical_backend import np, la
+
 SystemChain = Any # from oqupy.system import SystemChain
 
 class Gate:
@@ -59,7 +59,7 @@ def __init__(
 
         tmp_tensors = []
         for tensor in tensors:
-            tmp_tensor = np.array(tensor, dtype=NpDtype)
+            tmp_tensor = np.array(tensor, dtype=np.dtype_complex)
             tmp_tensors.append(tmp_tensor)
 
         for i, tensor in enumerate(tmp_tensors):
@@ -197,16 +197,16 @@ def compute_nn_gate(
         Nearest neighbor gate.
     """
     # exponentiate the liouvillian to become a propagator
-    propagator = linalg.expm(dt*liouvillian)
+    propagator = la.expm(dt*liouvillian)
     # split leg 0 and leg 1 each into left and right.
-    propagator.shape = [hs_dim_l**2, # left output
-                        hs_dim_r**2, # right output
-                        hs_dim_l**2, # left input
-                        hs_dim_r**2] # right input
+    propagator = propagator.reshape([hs_dim_l**2,   # left output
+                                     hs_dim_r**2,   # right output
+                                     hs_dim_l**2,   # left input
+                                     hs_dim_r**2])  # right input
     temp = np.swapaxes(propagator, 1, 2)
     temp = temp.reshape([hs_dim_l**2 * hs_dim_l**2,
                             hs_dim_r**2 * hs_dim_r**2])
-    u_full, s_full, vh_full = linalg.svd(temp, full_matrices=False)
+    u_full, s_full, vh_full = la.svd(temp, full_matrices=False)
     chi = _truncation_index(s_full, epsrel)
     s = s_full[:chi]
     u=u_full[:,:chi]
@@ -390,17 +390,20 @@ def __init__(
 
         tmp_gammas = []
         for g in gammas:
-            tmp_gamma = np.array(g, dtype=NpDtype)
+            tmp_gamma = np.array(g, dtype=np.dtype_complex)
             shape = deepcopy(tmp_gamma.shape)
             rank = len(shape)
             if rank == 4:
                 tmp_gamma = np.swapaxes(tmp_gamma,2,3)
             elif rank == 3:
-                tmp_gamma.shape = (shape[0], shape[1], 1, shape[2])
+                tmp_gamma = tmp_gamma.reshape((shape[0], shape[1], \
+                                               1, shape[2]))
             elif rank == 2:
-                tmp_gamma.shape = (1, shape[0]*shape[1], 1, 1)
+                tmp_gamma = tmp_gamma.reshape((1, shape[0] * shape[1], \
+                                               1, 1))
             elif rank == 1:
-                tmp_gamma.shape = (1, shape[0], 1, 1)
+                tmp_gamma = tmp_gamma.reshape((1, shape[0], \
+                                               1, 1))
             else:
                 raise ValueError()
             tmp_gammas.append(tmp_gamma)
@@ -416,9 +419,9 @@ def __init__(
         tmp_lambdas = []
         for bond_dim, l in zip(bond_dims, lambdas):
             if l is None:
-                tmp_lambda = np.ones(bond_dim, dtype=NpDtype)
+                tmp_lambda = np.ones(bond_dim, dtype=np.dtype_complex)
             else:
-                tmp_lambda = np.array(l, dtype=NpDtype)
+                tmp_lambda = np.array(l, dtype=np.dtype_complex)
                 shape = tmp_lambda.shape
                 rank = len(shape)
                 if rank == 2:
diff --git a/oqupy/operators.py b/oqupy/operators.py
index 88421af..71139e0 100644
--- a/oqupy/operators.py
+++ b/oqupy/operators.py
@@ -15,11 +15,9 @@
 
 from typing import Text
 
-import numpy as np
 from numpy import ndarray
 
-from oqupy.config import NpDtype
-
+from oqupy.backends.numerical_backend import np
 
 SIGMA = {"id":[[1, 0], [0, 1]],
          "x":[[0, 1], [1, 0]],
@@ -53,7 +51,7 @@ def identity(n: int) -> ndarray:
     identity : ndarray
         Identity matrix of dimension `n` x `n`.
     """
-    return np.identity(n, dtype=NpDtype)
+    return np.identity(n, dtype=np.dtype_complex)
 
 
 def sigma(name: Text) -> ndarray:
@@ -69,7 +67,7 @@ def sigma(name: Text) -> ndarray:
     sigma : ndarray
         Spin matrix of type `name`.
     """
-    return np.array(SIGMA[name], dtype=NpDtype)
+    return np.array(SIGMA[name], dtype=np.dtype_complex)
 
 
 def spin_dm(name: Text) -> ndarray:
@@ -86,7 +84,7 @@ def spin_dm(name: Text) -> ndarray:
     density_matrix : ndarray
         Spin density matrix.
     """
-    return np.array(SPIN_DM[name], dtype=NpDtype)
+    return np.array(SPIN_DM[name], dtype=np.dtype_complex)
 
 
 def create(n: int) -> ndarray:
@@ -120,7 +118,7 @@ def destroy(n: int) -> ndarray:
     create : ndarray
         Annihilation operator matrix of dimension `n` x `n`.
     """
-    return np.diag(np.sqrt(range(1, n), dtype=NpDtype), 1)
+    return np.diag(np.sqrt(np.arange(1, n, dtype=np.dtype_complex)), 1)
 
 # -- superoperators ----------------------------------------------------------
 
@@ -156,7 +154,7 @@ def preparation(
         density_matrix: ndarray) -> ndarray:
     """Construct the super operator that prepares the state. """
     dim = density_matrix.shape[0]
-    identity_matrix = np.identity(dim, dtype=NpDtype)
+    identity_matrix = np.identity(dim, dtype=np.dtype_complex)
     return np.outer(density_matrix.flatten(), identity_matrix.flatten())
 
 
diff --git a/oqupy/process_tensor.py b/oqupy/process_tensor.py
index d5df062..a600839 100644
--- a/oqupy/process_tensor.py
+++ b/oqupy/process_tensor.py
@@ -25,16 +25,16 @@
 from typing import Optional, Text, Union
 import warnings
 
-import numpy as np
+import h5py
 from numpy import ndarray
 import tensornetwork as tn
-import h5py
 
 from oqupy.base_api import BaseAPIClass
-from oqupy.config import NpDtype
-from oqupy import util
+from oqupy.util import create_delta
 from oqupy.version import __version__
 
+from oqupy.backends.numerical_backend import np
+
 class BaseProcessTensor(BaseAPIClass, ABC):
     """
     Abstract base class for process tensors in matrix product operator form
@@ -52,12 +52,12 @@ def __init__(
         self._hs_dim = hilbert_space_dimension
         self._dt = dt
         self._rho_dim = self._hs_dim**2
-        self._trace = (np.identity(self._hs_dim, dtype=NpDtype) \
+        self._trace = (np.identity(self._hs_dim, dtype=np.dtype_complex) \
                        / np.sqrt(float(self._hs_dim))).flatten()
         self._trace_square = self._trace**2
 
         if transform_in is not None:
-            tmp_transform_in = np.array(transform_in, dtype=NpDtype)
+            tmp_transform_in = np.array(transform_in, dtype=np.dtype_complex)
             assert len(tmp_transform_in.shape) == 2
             assert tmp_transform_in.shape[0] == self._rho_dim
             self._in_dim = tmp_transform_in.shape[1]
@@ -69,7 +69,7 @@ def __init__(
             self._trace_in = self._trace
 
         if transform_out is not None:
-            tmp_transform_out = np.array(transform_out, dtype=NpDtype)
+            tmp_transform_out = np.array(transform_out, dtype=np.dtype_complex)
             assert len(tmp_transform_out.shape) == 2
             assert tmp_transform_out.shape[1] == self._rho_dim
             self._out_dim = tmp_transform_out.shape[0]
@@ -240,7 +240,7 @@ def get_cap_tensor(self, step: int) -> ndarray:
         """
         Get the cap tensor (vector) to terminate the PT-MPO at time step `step`.
         """
-        return np.array([1.0], dtype=NpDtype)
+        return np.array([1.0], dtype=np.dtype_complex)
 
     def get_bond_dimensions(self) -> ndarray:
         """Return the bond dimensions of the MPS form of the process tensor."""
@@ -283,7 +283,8 @@ def set_initial_tensor(
         """
         if initial_tensor is None:
             self._initial_tensor = None
-            self._initial_tensor = np.array(initial_tensor, dtype=NpDtype)
+            self._initial_tensor = np.array(initial_tensor, \
+                                            dtype=np.dtype_complex)
 
     def set_mpo_tensor(
             self,
@@ -303,7 +304,7 @@ def set_mpo_tensor(
         length = len(self._mpo_tensors)
         if step >= length:
             self._mpo_tensors.extend([None] * (step - length + 1))
-        self._mpo_tensors[step] = np.array(tensor, dtype=NpDtype)
+        self._mpo_tensors[step] = np.array(tensor, dtype=np.dtype_complex)
 
     def set_cap_tensor(
             self,
@@ -315,7 +316,7 @@ def set_cap_tensor(
         length = len(self._cap_tensors)
         if step >= length:
             self._cap_tensors.extend([None] * (step - length + 1))
-        self._cap_tensors[step] = np.array(tensor, dtype=NpDtype)
+        self._cap_tensors[step] = np.array(tensor, dtype=np.dtype_complex)
 
     def get_initial_tensor(self) -> ndarray:
         """
@@ -344,7 +345,7 @@ def get_mpo_tensor(
             raise IndexError("Process tensor index out of bound. ")
         tensor = self._mpo_tensors[step]
         if len(tensor.shape) == 3:
-            tensor = util.create_delta(tensor, [0, 1, 2, 2])
+            tensor = create_delta(tensor, [0, 1, 2, 2])
         if transformed is False:
             return tensor
         if self._transform_in is not None:
@@ -383,7 +384,7 @@ def compute_caps(self) -> None:
         """
         length = len(self)
 
-        caps = [np.array([1.0], dtype=NpDtype)]
+        caps = [np.array([1.0], dtype=np.dtype_complex)]
         last_cap = tn.Node(caps[-1])
 
         for step in reversed(range(length)):
@@ -748,7 +749,7 @@ def get_mpo_tensor(
                                      shape=self._mpo_tensors_shape)
         if transformed:
             if len(tensor.shape) == 3:
-                tensor = util.create_delta(tensor, [0, 1, 2, 2])
+                tensor = create_delta(tensor, [0, 1, 2, 2])
             if self._transform_in is not None:
                 tensor = np.dot(np.moveaxis(tensor, -2, -1),
                                 self._transform_in.T)
@@ -783,7 +784,7 @@ def compute_caps(self) -> None:
         """
         length = len(self)
 
-        cap = np.array([1.0], dtype=NpDtype)
+        cap = np.array([1.0], dtype=np.dtype_complex)
         self.set_cap_tensor(length, cap)
         last_cap = tn.Node(cap)
 
diff --git a/oqupy/pt_tebd.py b/oqupy/pt_tebd.py
index 350b5fd..46d975f 100644
--- a/oqupy/pt_tebd.py
+++ b/oqupy/pt_tebd.py
@@ -21,22 +21,19 @@
 
 from typing import Dict, List, Optional, Text, Union
 
-import numpy as np
-
 from oqupy.backends.pt_tebd_backend import PtTebdBackend
 from oqupy.base_api import BaseAPIClass
-from oqupy.config import PT_TEBD_DEFAULT_ORDER
-from oqupy.config import PT_TEBD_DEFAULT_EPSREL
+from oqupy.config import PT_TEBD_DEFAULT_EPSREL, PT_TEBD_DEFAULT_ORDER
 from oqupy.control import ChainControl
 from oqupy.dynamics import Dynamics
-from oqupy.mps_mpo import GateLayer, SiteGate
-from oqupy.mps_mpo import compute_tebd_propagator
-from oqupy.mps_mpo import AugmentedMPS
-from oqupy.process_tensor import BaseProcessTensor
-from oqupy.process_tensor import TrivialProcessTensor
+from oqupy.mps_mpo import compute_tebd_propagator, AugmentedMPS, \
+    GateLayer, SiteGate
+from oqupy.process_tensor import BaseProcessTensor, TrivialProcessTensor
 from oqupy.system import SystemChain
 from oqupy.util import get_progress
 
+from oqupy.backends.numerical_backend import np
+
 NoneType = type(None)
 
 class PtTebdParameters(BaseAPIClass):
diff --git a/oqupy/pt_tempo.py b/oqupy/pt_tempo.py
index 9cc6f82..5aeb019 100644
--- a/oqupy/pt_tempo.py
+++ b/oqupy/pt_tempo.py
@@ -36,31 +36,26 @@
 200401 (2021).
 """
 
-from typing import Dict, Optional, Text, Union
 from copy import copy
+from typing import Dict, Optional, Text, Union
 
-import numpy as np
 from numpy import ndarray
 
+from oqupy.backends.pt_tempo_backend import PtTempoBackend
 from oqupy.base_api import BaseAPIClass
 from oqupy.bath import Bath
-from oqupy.config import PT_DEFAULT_TOLERANCE
-from oqupy.config import PT_TEMPO_BACKEND_CONFIG
-from oqupy.process_tensor import BaseProcessTensor
-from oqupy.process_tensor import SimpleProcessTensor
-from oqupy.process_tensor import FileProcessTensor
-from oqupy.backends.pt_tempo_backend import PtTempoBackend
-from oqupy.tempo import TempoParameters
-from oqupy.tempo import guess_tempo_parameters
-from oqupy.tempo import influence_matrix
+from oqupy.config import PT_DEFAULT_TOLERANCE, PT_TEMPO_BACKEND_CONFIG
 from oqupy.operators import left_right_super
+from oqupy.process_tensor import BaseProcessTensor, \
+    SimpleProcessTensor, FileProcessTensor
+from oqupy.tempo import guess_tempo_parameters, influence_matrix, \
+    TempoParameters
 from oqupy.util import get_progress
 
+from oqupy.backends.numerical_backend import np
 
 PT_CLASS = {"simple": SimpleProcessTensor}
 
-
-
 class PtTempo(BaseAPIClass):
     """
     Class to facilitate a PT-TEMPO computation.
@@ -155,7 +150,7 @@ def __init__(
         self._init_pt_tempo_backend()
 
     def _init_simple_process_tensor(self):
-        """ToDo. """
+        """TODO. """
         unitary = self._bath.unitary_transform
         if not np.allclose(unitary, np.identity(self._dimension)):
             transform_in = left_right_super(unitary.conjugate().T,
@@ -174,7 +169,7 @@ def _init_simple_process_tensor(self):
             description=self.description)
 
     def _init_file_process_tensor(self, filename, overwrite):
-        """ToDo. """
+        """TODO. """
         unitary = self._bath.unitary_transform
         if not np.allclose(unitary, np.identity(self._dimension)):
             transform_in = left_right_super(unitary.conjugate().T,
@@ -251,7 +246,7 @@ def _influence(self, dk: int) -> ndarray:
             deg_positions=tmp_deg_positions)
 
     @property
-    def dimension(self) -> np.ndarray:
+    def dimension(self) -> ndarray:
         """Hilbert space dimension. """
         return copy(self._dimension)
 
diff --git a/oqupy/system.py b/oqupy/system.py
index 13cc2d2..195328d 100644
--- a/oqupy/system.py
+++ b/oqupy/system.py
@@ -13,21 +13,19 @@
 Module on physical information of the system.
 """
 
-from typing import Callable, List, Optional, Text, Tuple
-from inspect import getfullargspec
 from copy import copy
 from functools import lru_cache
+from inspect import getfullargspec
+from typing import Callable, List, Optional, Text, Tuple
 
-import numpy as np
 from numpy import ndarray
-from scipy.linalg import expm
 from scipy import integrate
 from numdifftools import Jacobian
 
 from oqupy.base_api import BaseAPIClass
-from oqupy.config import NpDtype
 from oqupy import operators as opr
 
+from oqupy.backends.numerical_backend import np, la
 
 class BaseSystem(BaseAPIClass):
     """Base class for systems. """
@@ -120,8 +118,8 @@ def liouvillian(self) -> ndarray:
 
     def get_propagators(self, dt, start_time, subdiv_limit, epsrel):
         """Prepare propagator functions for the system. """
-        first_step = expm(self.liouvillian()*dt/2.0)
-        second_step = expm(self.liouvillian()*dt/2.0)
+        first_step = la.expm(self.liouvillian()*dt/2.0)
+        second_step = la.expm(self.liouvillian()*dt/2.0)
         def propagators(step: int):
             """Create the system propagators (first and second half) for
             the time step `step`  """
@@ -130,8 +128,8 @@ def propagators(step: int):
 
     def get_unitary_propagators(self, dt, start_time, subdiv_limit, epsrel):
         """Prepare propagator functions for the system. """
-        first_step = expm(-1j*self._hamiltonian*dt/2.0)
-        second_step = expm(-1j*self._hamiltonian*dt/2.0)
+        first_step = la.expm(-1j*self._hamiltonian*dt/2.0)
+        second_step = la.expm(-1j*self._hamiltonian*dt/2.0)
         def propagators(step: int):
             """Create the system propagators (first and second half) for
             the time step `step`  """
@@ -247,8 +245,8 @@ def propagators(step: int):
                 """Create the system propagators (first and second half) for
                 the time step `step`  """
                 t = start_time + step * dt
-                first_step = expm(self.liouvillian(t+dt/4.0)*dt/2.0)
-                second_step = expm(self.liouvillian(t+dt*3.0/4.0)*dt/2.0)
+                first_step = la.expm(self.liouvillian(t+dt/4.0)*dt/2.0)
+                second_step = la.expm(self.liouvillian(t+dt*3.0/4.0)*dt/2.0)
                 return first_step, second_step
         else:
             # Integrate Liouvillian to make propagators for first- and
@@ -257,12 +255,12 @@ def propagators(step: int):
                 """Create the system propagators (first and second half) for
                 the time step `step`  """
                 t = start_time + step * dt
-                first_step = expm(integrate.quad_vec(self.liouvillian,
+                first_step = la.expm(integrate.quad_vec(self.liouvillian,
                                                      a=t,
                                                      b=t+dt/2.0,
                                                      epsrel=epsrel,
                                                      limit=subdiv_limit)[0])
-                second_step = expm(integrate.quad_vec(self.liouvillian,
+                second_step = la.expm(integrate.quad_vec(self.liouvillian,
                                                       a=t+dt/2.0,
                                                       b=t+dt,
                                                       epsrel=epsrel,
@@ -440,9 +438,9 @@ def get_propagators(self, dt, start_time, subdiv_limit, epsrel):
             def propagators(step: int, field: complex,
                             field_derivative: complex):
                 t = start_time + step * dt
-                first_step = expm(self.liouvillian(t, t+dt/4.0,
+                first_step = la.expm(self.liouvillian(t, t+dt/4.0,
                     field, field_derivative)*dt/2.0)
-                second_step = expm(self.liouvillian(t, t+dt*3.0/4.0,
+                second_step = la.expm(self.liouvillian(t, t+dt*3.0/4.0,
                     field, field_derivative)*dt/2.0)
                 return first_step, second_step
         else:
@@ -453,12 +451,12 @@ def propagators(step: int, field: complex,
                 t = start_time + step * dt
                 liouvillian = lambda tau: self.liouvillian(t, tau,
                         field, field_derivative)
-                first_step = expm(integrate.quad_vec(liouvillian,
+                first_step = la.expm(integrate.quad_vec(liouvillian,
                                                      a=t,
                                                      b=t+dt/2.0,
                                                      epsrel=epsrel,
                                                      limit=subdiv_limit)[0])
-                second_step = expm(integrate.quad_vec(liouvillian,
+                second_step = la.expm(integrate.quad_vec(liouvillian,
                                                       a=t+dt/2.0,
                                                       b=t+dt,
                                                       epsrel=epsrel,
@@ -527,7 +525,8 @@ def __init__(
         """Create a ParameterizedSystem object."""
         # input check for Hamiltonian.
         number_of_parameters = len(getfullargspec(hamiltonian).args)
-        self._hamiltonian = np.vectorize(hamiltonian)
+        self._hamiltonian = np.vectorize(hamiltonian, \
+            signature=','.join(['()'] * number_of_parameters) + '->(m,m)')
         trial_hamiltonian = hamiltonian(*(list([0.5]*number_of_parameters)))
         _check_hamiltonian(trial_hamiltonian)
         dimension = trial_hamiltonian.shape[0]
@@ -556,7 +555,7 @@ def get_propagators(
             dt: float,
             parameters: ndarray) -> Callable[[int], Tuple[ndarray,ndarray]]:
         """
-        ToDo
+        TODO
         """
         def propagators(step: int):
             """Create the system propagators (first and second half) for
@@ -564,8 +563,8 @@ def propagators(step: int):
 
             pre_liou=self.liouvillian(*(list(parameters[2*step][:])))
             post_liou=self.liouvillian(*(list(parameters[2*step+1][:])))
-            first_step = expm(pre_liou*dt/2.0)
-            second_step = expm(post_liou*dt/2.0)
+            first_step = la.expm(pre_liou*dt/2.0)
+            second_step = la.expm(post_liou*dt/2.0)
 
             return first_step, second_step
         return propagators
@@ -579,13 +578,15 @@ def halfstep_propagator_derivative(self,dt):
         """
 
         def prop(parameterlist):
-            return expm(self.liouvillian(*parameterlist)*dt/2.0)
+            return la.expm(self.liouvillian(*parameterlist)*dt/2.0)
 
         jacfunre=Jacobian(lambda x: prop(x).real)
         jacfunim=Jacobian(lambda x: prop(x).imag)
 
         def jacfun(x):
-            jac=jacfunre(x)+1.0j*jacfunim(x)
+            # `npdifftools` returns a vanilla NumPy array
+            # TODO: use JAX-based differentiation modules
+            jac=np.array(jacfunre(x) + 1.0j * jacfunim(x))
 
             return [jac[:,i,:] for i in range(self._number_of_parameters)]
 
@@ -596,7 +597,7 @@ def get_propagator_derivatives(
             dt: float,
             parameters: ndarray) -> Callable[[int],Tuple[ndarray,ndarray]]:
         """
-        ToDo
+        TODO
         """
         if self._propagator_derivatives is not None:
             def propagator_derivatives_a(step: int):
@@ -718,13 +719,13 @@ def __init__(
         self._site_liouvillians = []
         for hs_dim in self._hs_dims:
             self._site_liouvillians.append(
-                np.zeros((hs_dim**2, hs_dim**2), dtype=NpDtype))
+                np.zeros((hs_dim**2, hs_dim**2), dtype=np.dtype_complex))
 
         self._nn_liouvillians = []
         for hs_dim_l, hs_dim_r in zip(self._hs_dims[:-1], self._hs_dims[1:]):
             self._nn_liouvillians.append(
                 np.zeros((hs_dim_l**2 * hs_dim_r**2, hs_dim_l**2 * hs_dim_r**2),
-                dtype=NpDtype))
+                dtype=np.dtype_complex))
 
         super().__init__(name, description)
 
@@ -770,7 +771,7 @@ def add_site_hamiltonian(
         assert isinstance(site, int)
         assert site >= 0
         assert site < len(self)
-        op = np.array(hamiltonian, dtype=NpDtype)
+        op = np.array(hamiltonian, dtype=np.dtype_complex)
         assert len(op.shape) == 2
         assert op.shape[0] == op.shape[1]
         assert self._hs_dims[site] == op.shape[0]
@@ -791,7 +792,8 @@ def add_site_liouvillian(
         liouvillian: ndarray
             Liouvillian acting on the single site.
         """
-        self._site_liouvillians[site] += np.array(liouvillian, dtype=NpDtype)
+        self._site_liouvillians[site] += np.array(liouvillian, \
+                                                  dtype=np.dtype_complex)
 
     def add_site_dissipation(
             self,
@@ -820,7 +822,7 @@ def add_site_dissipation(
         gamma: float
             Optional multiplicative factor :math:`\gamma`.
         """
-        op = np.array(lindblad_operator, dtype=NpDtype)
+        op = np.array(lindblad_operator, dtype=np.dtype_complex)
         op_dagger = op.conjugate().T
         self._site_liouvillians[site] += \
             gamma * (opr.left_right_super(op, op_dagger) \
@@ -855,8 +857,8 @@ def add_nn_hamiltonian(
         assert isinstance(site, int)
         assert site >= 0
         assert site < len(self) - 1
-        op_l = np.array(hamiltonian_l, dtype=NpDtype)
-        op_r = np.array(hamiltonian_r, dtype=NpDtype)
+        op_l = np.array(hamiltonian_l, dtype=np.dtype_complex)
+        op_r = np.array(hamiltonian_r, dtype=np.dtype_complex)
         assert len(op_l.shape) == 2
         assert len(op_r.shape) == 2
         assert op_l.shape[0] == op_l.shape[1]
@@ -881,7 +883,8 @@ def add_nn_liouvillian(
         liouvillian_l_r: ndarray
             Liouvillian acting on sites :math:`n` and :math:`n+1`.
         """
-        self._nn_liouvillians[site] += np.array(liouvillian_l_r, dtype=NpDtype)
+        self._nn_liouvillians[site] += np.array(liouvillian_l_r, \
+                                                dtype=np.dtype_complex)
 
     def add_nn_dissipation(
             self,
@@ -917,8 +920,8 @@ def add_nn_dissipation(
         assert isinstance(site, int)
         assert site >= 0
         assert site < len(self) - 1
-        op_l = np.array(lindblad_operator_l, dtype=NpDtype)
-        op_r = np.array(lindblad_operator_r, dtype=NpDtype)
+        op_l = np.array(lindblad_operator_l, dtype=np.dtype_complex)
+        op_r = np.array(lindblad_operator_r, dtype=np.dtype_complex)
         assert len(op_l.shape) == 2
         assert len(op_r.shape) == 2
         assert op_l.shape[0] == op_l.shape[1]
@@ -970,8 +973,10 @@ def get_nn_full_liouvillians(self) -> List[ndarray]:
 def _check_hamiltonian(hamiltonian) -> ndarray:
     """Input checking for a single Hamiltonian. """
     try:
-        tmp_hamiltonian = np.array(hamiltonian, dtype=NpDtype)
-        tmp_hamiltonian.setflags(write=False)
+        tmp_hamiltonian = np.array(hamiltonian, dtype=np.dtype_complex)
+        # set immutable if default NumPy array
+        if isinstance(tmp_hamiltonian, ndarray):
+            tmp_hamiltonian.setflags(write=False)
     except Exception as e:
         raise AssertionError("Coupling operator must be numpy array") from e
     assert len(tmp_hamiltonian.shape) == 2, \
@@ -985,7 +990,8 @@ def _check_tdependent_hamiltonian(hamiltonian) -> Callable[[float],
         ndarray]:
     """Input checking for a time-dependent Hamiltonian. """
     try:
-        tmp_hamiltonian = np.vectorize(hamiltonian)
+        tmp_hamiltonian = np.vectorize(hamiltonian, \
+            signature='()->(m,m)')
         _check_hamiltonian(tmp_hamiltonian(1.0))
     except Exception as e:
         raise AssertionError(
@@ -996,7 +1002,8 @@ def _check_tdependent_hamiltonian(hamiltonian) -> Callable[[float],
 def _check_tfielddependent_hamiltonian(hamiltonian) -> Callable[[float,
     complex], ndarray]:
     try:
-        tmp_hamiltonian = np.vectorize(hamiltonian)
+        tmp_hamiltonian = np.vectorize(hamiltonian, \
+            signature='(),()->(m,m)')
         _check_hamiltonian(tmp_hamiltonian(1.0, 1.0+1.0j))
     except Exception as e:
         raise AssertionError(
@@ -1035,7 +1042,7 @@ def _check_gammas_lindblad_operators(gammas, lindblad_operators) -> Tuple[
         tmp_lindblad_operators = []
         for lindblad_operator in lindblad_operators:
             tmp_lindblad_operators.append(
-                np.array(lindblad_operator, dtype=NpDtype))
+                np.array(lindblad_operator, dtype=np.dtype_complex))
     except Exception as e:
         raise AssertionError(
             "All elements of `lindblad_operators` must be numpy arrays.") \
@@ -1055,7 +1062,8 @@ def _check_tdependent_gammas_lindblad_operators(
         tmp_gammas = []
         for gamma in gammas:
             float(gamma(1.0))
-            tmp_gamma = np.vectorize(gamma)
+            tmp_gamma = np.vectorize(gamma, \
+                signature='()->()')
             tmp_gammas.append(tmp_gamma)
     except Exception as e:
         raise AssertionError(
@@ -1064,7 +1072,8 @@ def _check_tdependent_gammas_lindblad_operators(
     try:
         tmp_lindblad_operators = []
         for lindblad_operator in lindblad_operators:
-            tmp_lindblad_operator = np.vectorize(lindblad_operator)
+            tmp_lindblad_operator = np.vectorize(lindblad_operator, \
+                signature='()->(m,m)')
             np.array(tmp_lindblad_operator(1.0))
             tmp_lindblad_operators.append(tmp_lindblad_operator)
     except Exception as e:
@@ -1135,8 +1144,8 @@ def _imaginary_liouvillian(hamiltonian, gammas, lindblad_operators):
 def _create_density_matrix(dim, seed=1):
     r"""Create a repeatable (dim,dim) matrix that represents
     a valid density matrix :math:`rho`"""
-    rng = np.random.default_rng(seed)
-    a = rng.random((dim, dim)) + 1j*rng.random((dim,dim))
+    random_floats = np.get_random_floats(seed, (2, dim, dim))
+    a = random_floats[0] + 1j * random_floats[1]
     b = np.matmul(a, a.conj().T)
     rho = b / b.trace()
     return rho
diff --git a/oqupy/system_dynamics.py b/oqupy/system_dynamics.py
index 1ddc965..c545666 100644
--- a/oqupy/system_dynamics.py
+++ b/oqupy/system_dynamics.py
@@ -17,22 +17,20 @@
 import warnings
 
 from itertools import product
-import numpy as np
 from numpy import ndarray
 import tensornetwork as tn
-from oqupy.system import TimeDependentSystemWithField
 
-from oqupy.config import NpDtype, INTEGRATE_EPSREL, SUBDIV_LIMIT
+from oqupy.config import INTEGRATE_EPSREL, SUBDIV_LIMIT
 from oqupy.control import Control
 from oqupy.dynamics import Dynamics, MeanFieldDynamics
-from oqupy.process_tensor import BaseProcessTensor
-from oqupy.system import BaseSystem, System, TimeDependentSystem
-from oqupy.system import ParameterizedSystem
-from oqupy.system import MeanFieldSystem
 from oqupy.operators import left_super, right_super
-from oqupy.util import check_convert, check_isinstance, check_true
-from oqupy.util import get_progress
+from oqupy.process_tensor import BaseProcessTensor
+from oqupy.system import BaseSystem, System, TimeDependentSystem, \
+    TimeDependentSystemWithField, ParameterizedSystem, MeanFieldSystem
+from oqupy.util import check_convert, check_isinstance, check_true, \
+    get_progress
 
+from oqupy.backends.numerical_backend import np
 
 Indices = Union[int, slice, List[Union[int, slice]]]
 
@@ -119,7 +117,7 @@ def controls(step: int):
     #    edges 0, 1, .., num_envs-1    are the bond legs of the environments
     #    edge  -1                      is the state leg
     initial_ndarray = initial_state.reshape(hs_dim**2)
-    initial_ndarray.shape = tuple([1]*num_envs+[hs_dim**2])
+    initial_ndarray = initial_ndarray.reshape(tuple([1]*num_envs+[hs_dim**2]))
     current_node = tn.Node(initial_ndarray)
     current_edges = current_node[:]
 
@@ -344,7 +342,8 @@ def prepare_controls(step: int, control:Control):
                parsed_parameters_dict["hs_dim"],
                num_envs_list):
         initial_ndarray = initial_state.reshape(hs_dim**2)
-        initial_ndarray.shape = tuple([1]*num_envs+[hs_dim**2])
+        initial_ndarray = initial_ndarray.reshape(tuple([1] * num_envs + \
+                                                        [hs_dim**2]))
         current_node = tn.Node(initial_ndarray)
         current_edges = current_node[:]
 
@@ -508,7 +507,7 @@ def _compute_dynamics_input_parse(
 
     if initial_state is None:
         raise ValueError("An initial state must be specified.")
-    check_isinstance(initial_state, ndarray, "initial_state")
+    check_isinstance(initial_state, np.ndarray, "initial_state")
     check_true(
         initial_state.shape == (hs_dim, hs_dim),
         "Initial sate must be a square matrix of " \
@@ -535,7 +534,7 @@ def _compute_dynamics_input_parse(
                     "All process tensors must have the same "\
                             "timestep length.")
         max_steps.append(pt.max_step)
-    max_step = np.min(max_steps+[np.inf])
+    max_step = np.min(np.array(max_steps + [np.inf]))
 
     if dt is None:
         raise ValueError(
@@ -560,7 +559,7 @@ def _compute_dynamics_input_parse(
 
 
 def _get_caps(process_tensors: List[BaseProcessTensor], step: int):
-    """ToDo """
+    """TODO """
     caps = []
     for i in range(len(process_tensors)):
         try:
@@ -577,7 +576,7 @@ def _get_caps(process_tensors: List[BaseProcessTensor], step: int):
 
 
 def _get_pt_mpos(process_tensors: List[BaseProcessTensor], step: int):
-    """ToDo """
+    """TODO """
     pt_mpos = []
     for i in range(len(process_tensors)):
         pt_mpo = process_tensors[i].get_mpo_tensor(step)
@@ -629,7 +628,7 @@ def _get_pt_mpos_backprop(mpo_list:ndarray, step: int):
 
 
 def _apply_system_superoperator(current_node, current_edges, sup_op):
-    """ToDo """
+    """TODO """
     if sup_op is None:
         return current_node, current_edges
     sup_op_node = tn.Node(sup_op.T)
@@ -641,10 +640,10 @@ def _apply_system_superoperator(current_node, current_edges, sup_op):
 
 
 def _apply_caps(current_node, current_edges, caps):
-    """ToDo """
+    """TODO """
     node_dict, edge_dict = tn.copy([current_node])
     for current_edge, cap in zip(current_edges[:-1], caps):
-        cap_node = tn.Node(cap)
+        cap_node = tn.Node(np.array(cap))
         edge_dict[current_edge] ^ cap_node[0]
         node_dict[current_node] = node_dict[current_node] @ cap_node
     state_node = node_dict[current_node]
@@ -889,9 +888,13 @@ def compute_correlations_nt(
         lengths = len(ops_times_[i])
         times_length.append(lengths)
 
-    ret_correlations = np.empty(times_length, dtype=NpDtype)
+    ret_correlations = np.empty(times_length, dtype=np.dtype_complex)
     #This array will contain all correlations
-    ret_correlations[:] = np.nan + 1.0j*np.nan
+    ret_correlations = np.update(
+        array=ret_correlations,
+        indices=(slice(None), ),
+        values=np.nan + 1.0j * np.nan
+    )
 
 
     parameters = {
@@ -918,7 +921,7 @@ def compute_correlations_nt(
             #check time ordering
             ft = np.array(first_times)
             check = sorted(first_times)
-            if not np.allclose(ft, check):
+            if not np.allclose(ft, np.array(check)):
                 continue
             ft_max = ft.max()
             if (ft_max > last_times).any():
@@ -935,7 +938,11 @@ def compute_correlations_nt(
                                                     operators = operators,
                                                     ops_order = ops_order,
                                                     **parameters)
-            ret_correlations[sch_indices[i]] = corr
+            ret_correlations = np.update(
+                array=ret_correlations,
+                indices=sch_indices[i],
+                values=corr
+            )
         prog_bar.update(len(schedule))
     return ret_times, ret_correlations
 
diff --git a/oqupy/tempo.py b/oqupy/tempo.py
index 896679b..0347869 100644
--- a/oqupy/tempo.py
+++ b/oqupy/tempo.py
@@ -24,29 +24,30 @@
 time-evolving matrix product operators*, Nat. Commun. 9, 3322 (2018).
 """
 
+from copy import copy
 import sys
 from typing import Callable, Dict, List, Optional, Text, Tuple, Union
 import warnings
-from copy import copy
 
-import numpy as np
 from numpy import ndarray
 
 from oqupy.bath import Bath
 from oqupy.base_api import BaseAPIClass
-from oqupy.config import MAX_DKMAX, DEFAULT_TOLERANCE, MAX_SYS_SAMPLES
+from oqupy.config import DEFAULT_TOLERANCE, MAX_DKMAX, MAX_SYS_SAMPLES
 from oqupy.config import INTEGRATE_EPSREL, SUBDIV_LIMIT
 from oqupy.config import TEMPO_BACKEND_CONFIG
 from oqupy.bath_correlations import BaseCorrelations, CustomSD
 from oqupy.dynamics import Dynamics, MeanFieldDynamics
-from oqupy.system import BaseSystem, System, TimeDependentSystem,\
+from oqupy.system import BaseSystem, System, TimeDependentSystem, \
     TimeDependentSystemWithField, MeanFieldSystem
 from oqupy.backends.tempo_backend import TempoBackend
 from oqupy.backends.tempo_backend import MeanFieldTempoBackend
 from oqupy.backends.tempo_backend import TIBaseBackend
-from oqupy.util import check_convert, check_isinstance, check_true,\
+from oqupy.util import check_convert, check_isinstance, check_true, \
         get_progress
 
+from oqupy.backends.numerical_backend import np
+
 class TempoParameters(BaseAPIClass):
     r"""
     Parameters for the TEMPO computation.
@@ -578,7 +579,7 @@ def coeffs(k):
 
 
         unitary_transform = self._bath.unitary_transform
-        # ToDo: Unitary transform is not used. Check that GibbsTempo also works
+        # TODO: Unitary transform is not used. Check that GibbsTempo also works
         #       with non-diagonal coupling operators!!!
         if not np.allclose(unitary_transform,np.identity(self.dimension)):
             raise NotImplementedError(
@@ -985,8 +986,8 @@ def influence_matrix(
         time_1 = float(dkmax) * dt
         if parameters.add_correlation_time is not None:
             time_2 = float(dkmax) * dt \
-                + np.min([float(-dk) * dt,
-                            1.0*dt + parameters.add_correlation_time])
+                + float(np.min(np.array([float(-dk) * dt, 1.0 * dt + \
+                                         parameters.add_correlation_time])))
         else:
             return None
         shape = "rectangle"
@@ -1286,7 +1287,7 @@ def _tempo_physical_input_parse(
     hs_dim = system.dimension
 
     if initial_state is not None:  # initial state is None for Gibbs Tempo
-        check_isinstance(initial_state, ndarray, "initial_state")
+        check_isinstance(initial_state, np.ndarray, "initial_state")
         check_true(
             initial_state.shape == (hs_dim, hs_dim),
             "Initial sate must be a square matrix of " \
@@ -1334,9 +1335,9 @@ def _estimate_dt_dkmax_from_bath(bath, start_time, end_time, tolerance):
     return dt, dkmax
 
 def _analyse_correlation(
-        corr_func: Callable[[np.ndarray],np.ndarray],
-        times: np.ndarray,
-        corr_vals: np.ndarray):
+        corr_func: Callable[[ndarray],ndarray],
+        times: ndarray,
+        corr_vals: ndarray):
     """Check correlation function on a finer grid."""
     additional_times = (times[:-1] + times[1:])/2.0
     additional_corr_vals = corr_func(additional_times)
diff --git a/oqupy/util.py b/oqupy/util.py
index 07d8c91..b90a2df 100644
--- a/oqupy/util.py
+++ b/oqupy/util.py
@@ -13,75 +13,74 @@
 Module for utilities.
 """
 
+from copy import copy
+from datetime import timedelta
 import sys
-import copy as cp
-from typing import Any, List, Optional, Text
 from threading import Timer
 from time import time
-from datetime import timedelta
+from typing import Any, List, Optional, Text
 
-import numpy as np
 from numpy import ndarray
 
 from oqupy.config import PROGRESS_TYPE
 
+from oqupy.backends.numerical_backend import np
+
 # -- numpy utils --------------------------------------------------------------
 
 def create_delta(
         tensor: ndarray,
-        index_scrambling: List[int]) -> ndarray:
-    """
-    Creates deltas in numpy tensor.
-
-    .. warning::
-        This is a computationally inefficient method to perform the task.
-
-    .. todo::
-        Make it better.
-
-    """
-    tensor_shape = tensor.shape
-    a = [0]*len(tensor_shape)
-
-    ret_shape = tuple(list(tensor_shape)[i] for i in index_scrambling)
-    ret_ndarray = np.zeros(ret_shape, dtype=tensor.dtype)
-
-    # emulating a do-while loop
-    do_while_condition = True
-    while do_while_condition:
-        tensor_indices = tuple(a)
-        ret_indices = tuple(a[i] for i in index_scrambling)
-        ret_ndarray[ret_indices] = tensor[tensor_indices]
-        do_while_condition = increase_list_of_index(a, tensor_shape)
-
-    return ret_ndarray
-
-def increase_list_of_index(
-        a: List,
-        shape: List,
-        index: Optional[int] = -1) -> bool:
-    """Circle through a list of indices. """
-    a[index] += 1
-    if a[index] >= shape[index]:
-        if index == -len(shape):
-            return False
-        a[index] = 0
-        return increase_list_of_index(a, shape, index-1)
-    return True
+        index_scrambling: List[int]
+    ) -> ndarray:
+    """Creates deltas in a tensor."""
+    # converting to NumPy-array for future-proof implementation
+    # see [this issue](https://github.com/google/jax/issues/4564)
+    # the shape of the tensor has n_in elements whereas
+    # index_scrambling has n_out elements
+    _shape = np.array(tensor.shape, dtype=int)
+    _idxs = np.array(index_scrambling, dtype=int)
+
+    # obtain the selection indices for each axis
+    _indices = get_indices(_shape, np.prod(_shape))
+
+    # scramble output tensor with elements of input tensor
+    return np.update(
+        array=np.zeros(tuple(_shape[_idxs]),
+                       dtype=tensor.dtype),
+        indices=tuple(_indices[_idxs]),
+        values=tensor[tuple(_indices)]
+    )
+
+def get_indices(
+        shape: ndarray,
+        n_iters: int,
+    ) -> ndarray:
+    """Obtain index matrix for scrambling."""
+    # obtain divisors for each axis as values equal to the
+    # number of elements contained upto the preceeding axes
+    # for e.g., shape [4, 5, 3] will result in [15, 3, 1]
+    divisors = np.cumprod(np.concatenate([
+        shape[1:],
+        np.array([1], dtype=int)
+    ])[::-1])[::-1]
+    # prepare an iteration matrix of shape (n_iters, n_in)
+    # to index each axis, for e.g., n_iters = 3 x 5 x 4
+    iteration_matrix = np.arange(0, n_iters).reshape(
+        (n_iters, 1)).repeat(shape.shape[0], 1)
+    # divide each element with the divisors obtained above
+    # and obtain the remainder modullo the size of each axis
+    # return the index matrix with shape (n_in, n_iters)
+    return ((iteration_matrix / divisors).astype(int) % shape).T
 
 def add_singleton(
         tensor: ndarray,
-        index: Optional[int] = -1,
-        copy: Optional[bool] = True) -> ndarray:
+        index: Optional[int] = -1) -> ndarray:
     """Add a singleton to a numpy tensor. """
-    if copy:
-        ten = cp.copy(tensor)
-    else:
-        ten = tensor
-    shape_list = list(ten.shape)
+    # TODO: Check if the array is being copied twice
+    tensor = copy(tensor)
+    shape_list = list(tensor.shape)
     shape_list.insert(index, 1)
-    ten.shape = tuple(shape_list)
-    return ten
+    return tensor.reshape(shape_list)
 
 def is_diagonal_matrix(tensor: ndarray):
     """Check if matrix is diagonal """
diff --git a/pylintrc b/pylintrc
index 2a2ec1e..65e1959 100644
--- a/pylintrc
+++ b/pylintrc
@@ -79,6 +79,7 @@ disable=raw-checker-failed,
         consider-using-f-string,          # GEFux: added for development
         too-many-arguments,               # piperfw: added for development
         too-many-positional-arguments,    # piperfw: added 2024-09-27
+        import-outside-toplevel,          # sampreet: added 2024-10-24 for JAX
         possibly-used-before-assignment,  # GEFux: added by hand
         unnecessary-lambda-assignment     # GEFux: added by hand
 


Test patch:
diff --git a/tests/coverage/__init__.py b/tests/coverage/__init__.py
index e69de29..9ae2946 100644
--- a/tests/coverage/__init__.py
+++ b/tests/coverage/__init__.py
@@ -0,0 +1,2 @@
+# from oqupy.backends import enable_jax_features
+# enable_jax_features()
diff --git a/tests/coverage/bath_test.py b/tests/coverage/bath_test.py
index b4d6fc4..5e5a706 100644
--- a/tests/coverage/bath_test.py
+++ b/tests/coverage/bath_test.py
@@ -15,11 +15,10 @@
 
 import pytest
 
-import numpy as np
-
 from oqupy.bath import Bath
 from oqupy.bath_correlations import PowerLawSD
-from oqupy import operators
+
+from oqupy.backends.numerical_backend import default_np, np
 
 def test_bath():
     wc = 4.0
@@ -43,7 +42,7 @@ def test_bath():
 
     # see if the properties work
     str(bath_A)
-    np.testing.assert_equal(bath_A.coupling_operator, coupling_operator)
+    default_np.testing.assert_array_equal(bath_A.coupling_operator, coupling_operator)
     assert bath_A.dimension == 2
     assert bath_A.correlations.zeta == 1.0
     assert max(bath_A.west_degeneracy_map) == 2
@@ -76,8 +75,8 @@ def test_bath():
     # try 1 dimensional coupling degeneracies
     coupling_op = np.array([[1.0,0.0],[0.0,1.0]])
     bath_C = Bath(coupling_op, correlations)
-    assert np.array_equal(bath_C.west_degeneracy_map,[0,0,0,0])
-    assert np.array_equal(bath_C.north_degeneracy_map,[0,0,0,0])
+    default_np.testing.assert_array_equal(bath_C.west_degeneracy_map,[0,0,0,0])
+    default_np.testing.assert_array_equal(bath_C.north_degeneracy_map,[0,0,0,0])
 
 def test_bath_bad_input():
     wc = 4.0
diff --git a/tests/coverage/contractions_test.py b/tests/coverage/contractions_test.py
index 1445256..2dc06aa 100644
--- a/tests/coverage/contractions_test.py
+++ b/tests/coverage/contractions_test.py
@@ -14,11 +14,10 @@
 """
 
 import pytest
-import numpy as np
 
 import oqupy
-from oqupy.system_dynamics import compute_correlations_nt
 
+from oqupy.backends.numerical_backend import np
 
 def test_compute_dynamics_with_field():
     start_time = 1
@@ -56,7 +55,7 @@ def test_compute_dynamics_with_field():
     assert isinstance(mean_field_dynamics, oqupy.dynamics. MeanFieldDynamics)
     assert len(mean_field_dynamics.times) == 1 + num_steps
     assert np.isclose(np.min(mean_field_dynamics.times), start_time)
-    assert type(mean_field_dynamics.fields[1]) == oqupy.config.NpDtype
+    assert mean_field_dynamics.fields[1].dtype == oqupy.config.NumPyDtypeComplex
     assert np.isclose(mean_field_dynamics.fields[0], initial_field)
     # check correct type of dynamics returned
     assert isinstance(mean_field_dynamics, oqupy.dynamics.MeanFieldDynamics)
@@ -65,7 +64,7 @@ def test_compute_dynamics_with_field():
     # check start_time used correctly
     assert np.isclose(np.min(mean_field_dynamics.times), start_time)
     # check complex type
-    assert type(mean_field_dynamics.fields[1]) == oqupy.config.NpDtype
+    assert mean_field_dynamics.fields[1].dtype == oqupy.config.NumPyDtypeComplex
     # check initial field recorded correctly
     assert np.isclose(mean_field_dynamics.fields[0], initial_field)
 
@@ -236,7 +235,7 @@ def test_compute_correlations_nt():
                                             end_time= dt * 10,
                                             parameters=tempo_parameters)
 
-    cor = compute_correlations_nt(system = system,
+    cor = oqupy.compute_correlations_nt(system = system,
                                       process_tensor=process_tensor,
                                       operators = operators,
                                       ops_times=ops_times,
@@ -250,12 +249,12 @@ def test_compute_correlations_nt():
     for i in range (len(ops_times)):
         assert len(cor[0][i]) == cor[1].shape[i]
     assert len(operators) == len(ops_times) == len(time_order)
-    assert type(cor[1][0][0][0][0]) == oqupy.config.NpDtype
+    assert cor[1][0][0][0][0].dtype == oqupy.config.NumPyDtypeComplex
 
     # input checks
     # no process tensor
     with pytest.raises(AssertionError):
-        compute_correlations_nt(system = system,
+        oqupy.compute_correlations_nt(system = system,
                                           process_tensor=None,
                                           operators = operators,
                                           ops_times=ops_times,
@@ -266,7 +265,7 @@ def test_compute_correlations_nt():
                                           progress_type = "bar")
     # time arguments and no. of operators don't match
     with pytest.raises(AssertionError):
-        compute_correlations_nt(system = system,
+        oqupy.compute_correlations_nt(system = system,
                                           process_tensor=process_tensor,
                                           operators = [sigma_x,
                                                         sigma_x, sigma_x],
@@ -278,7 +277,7 @@ def test_compute_correlations_nt():
                                           progress_type = "bar")
     # time argument exceeds process tensor length
     with pytest.raises(IndexError):
-        compute_correlations_nt(system = system,
+        oqupy.compute_correlations_nt(system = system,
                                           process_tensor=process_tensor,
                                           operators = operators,
                                           ops_times=[0*dt, 1*dt, 2*dt, 12*dt],
@@ -288,7 +287,7 @@ def test_compute_correlations_nt():
                                           start_time = start_time,
                                           progress_type = "bar")
     #no specified time step
-    cor=compute_correlations_nt(system = system,
+    cor=oqupy.compute_correlations_nt(system = system,
                                           process_tensor=process_tensor,
                                           operators = operators,
                                           ops_times=ops_times,
@@ -300,7 +299,7 @@ def test_compute_correlations_nt():
     assert cor[0][0][1]-cor[0][0][0] == process_tensor.dt
     #time out of bound
     with pytest.raises(IndexError):
-        cor=compute_correlations_nt(system = system,
+        cor=oqupy.compute_correlations_nt(system = system,
                                           process_tensor=process_tensor,
                                           operators = operators,
                                           ops_times=[0, 1, 2, 12],
@@ -311,7 +310,7 @@ def test_compute_correlations_nt():
                                           progress_type = "bar")
     #Wrong time argument type
     with pytest.raises(TypeError):
-        cor=compute_correlations_nt(system = system,
+        cor=oqupy.compute_correlations_nt(system = system,
                                           process_tensor=process_tensor,
                                           operators = operators,
                                           ops_times=["blah", 1.*dt, 2.*dt,
@@ -347,7 +346,7 @@ def test_compute_correlations_nt_B():
                                             end_time= dt * 10,
                                             parameters=tempo_parameters)
 
-    cor = compute_correlations_nt(system = system,
+    cor = oqupy.compute_correlations_nt(system = system,
                                       process_tensor=process_tensor,
                                       operators = operators,
                                       ops_times=ops_times,
@@ -359,7 +358,7 @@ def test_compute_correlations_nt_B():
     assert np.isnan(cor[1][0][0])
 
     ops_times = [5*dt, (4*dt, 6*dt)]
-    cor = compute_correlations_nt(system = system,
+    cor = oqupy.compute_correlations_nt(system = system,
                                       process_tensor=process_tensor,
                                       operators = operators,
                                       ops_times=ops_times,
@@ -369,6 +368,4 @@ def test_compute_correlations_nt_B():
                                       start_time = start_time,
                                       progress_type = "bar")
     assert np.isnan(cor[1][0][0])
-    assert type(cor[1][0][1]) == oqupy.config.NpDtype
-
-
+    assert cor[1][0][1].dtype == oqupy.config.NumPyDtypeComplex
diff --git a/tests/coverage/correlations_test.py b/tests/coverage/correlations_test.py
index 6a67a53..f7def2b 100644
--- a/tests/coverage/correlations_test.py
+++ b/tests/coverage/correlations_test.py
@@ -14,12 +14,11 @@
 """
 
 import pytest
-import numpy as np
 
-from oqupy.bath_correlations import BaseCorrelations
-from oqupy.bath_correlations import CustomCorrelations
-from oqupy.bath_correlations import CustomSD
-from oqupy.bath_correlations import PowerLawSD
+from oqupy.bath_correlations import BaseCorrelations, CustomCorrelations, \
+    CustomSD, PowerLawSD
+
+from oqupy.backends.numerical_backend import default_np, np
 
 square_function = lambda w: 0.1 * w ** 2
 
@@ -96,11 +95,11 @@ def test_matsubara_custom_s_d():
         w = np.linspace(0, 8.0 * sd.cutoff, 10)
         y = sd.spectral_density(w)
         t = np.linspace(0, 4.0 / sd.cutoff, 10)
-        [sd.correlation(tt, matsubara=True) for tt in t]
-        [sd.eta_function(tt, matsubara=True) for tt in t]
+        [sd.correlation(float(tt), matsubara=True) for tt in t]
+        [sd.eta_function(float(tt), matsubara=True) for tt in t]
 
         assert type(sd.correlation(1, matsubara=True)) == float
-        np.testing.assert_almost_equal(
+        default_np.testing.assert_array_almost_equal(
             sd.correlation(0, matsubara=True),
             sd.correlation(1 / temperature, matsubara=True))
 
@@ -139,7 +138,7 @@ def test_matsubara_custom_s_d():
                                               matsubara=ma)
             # assert np.round(big_tri - (3 * small_tri + square + rect), 14) == 0
             stiched_big_tri = 3 * mid_tri + square + rect - 2*small_tri
-            np.testing.assert_almost_equal( stiched_big_tri, big_tri)
+            default_np.testing.assert_array_almost_equal(stiched_big_tri, big_tri)
 
 def test_matsubara_custom_s_d_bad_input():
     temperature = 0.0
diff --git a/tests/coverage/dynamics_test.py b/tests/coverage/dynamics_test.py
index e0f90ef..3a06669 100644
--- a/tests/coverage/dynamics_test.py
+++ b/tests/coverage/dynamics_test.py
@@ -14,10 +14,11 @@
 """
 
 import pytest
-import numpy as np
 
 from oqupy.dynamics import BaseDynamics, Dynamics, MeanFieldDynamics
 
+from oqupy.backends.numerical_backend import default_np, np
+
 
 times = [0.0, 0.1]
 states = [np.array([[1, 0], [0, 0]]), np.array([[0.9, -0.1j], [0.1j, 0.1]])]
@@ -30,8 +31,8 @@ def test_base_dynamics():
     dyn_A = BaseDynamics()
     str(dyn_A)
     len(dyn_A)
-    assert type(dyn_A.times) == np.ndarray
-    assert type(dyn_A.states) == np.ndarray
+    assert isinstance(dyn_A.times, np.ndarray)
+    assert isinstance(dyn_A.states, np.ndarray)
     assert dyn_A.shape is None
     t, expec = dyn_A.expectations()
     assert t is None
@@ -44,9 +45,9 @@ def test_dynamics():
     dyn_B = Dynamics(times, states)
     str(dyn_B)
     assert len(dyn_B) == len(times)
-    np.testing.assert_almost_equal(times, dyn_B.times)
-    np.testing.assert_almost_equal(states, dyn_B.states)
-    np.testing.assert_almost_equal(states[0].shape, dyn_B.shape)
+    default_np.testing.assert_array_almost_equal(times, dyn_B.times)
+    default_np.testing.assert_array_almost_equal(states, dyn_B.states)
+    default_np.testing.assert_array_almost_equal(states[0].shape, dyn_B.shape)
     dyn_B.add(other_time, other_state)
 
 def test_dynamics_bad_input():
@@ -62,12 +63,12 @@ def test_dynamics_bad_input():
     with pytest.raises(AssertionError):
         dyn_A.add(other_time, "bla")
     with pytest.raises(AssertionError):
-        dyn_A.add(other_time, np.random.rand(3,3))
+        dyn_A.add(other_time, np.get_random_floats(0, (3,3)))
     dyn_B = Dynamics()
     with pytest.raises(AssertionError):
-        dyn_A.add(other_time, np.random.rand(2,2,2))
+        dyn_A.add(other_time, np.get_random_floats(1, (2,2,2)))
     with pytest.raises(AssertionError):
-        dyn_A.add(other_time, np.random.rand(3,2))
+        dyn_A.add(other_time, np.get_random_floats(2, (3,2)))
 
 def test_dynamics_expectations():
     dyn_A = Dynamics()
@@ -76,14 +77,14 @@ def test_dynamics_expectations():
     assert x is None
     dyn = Dynamics(times, states)
     t, tr = dyn.expectations()
-    np.testing.assert_almost_equal(t, dyn.times)
-    np.testing.assert_almost_equal(tr, [1.0, 1.0])
+    default_np.testing.assert_array_almost_equal(t, dyn.times)
+    default_np.testing.assert_array_almost_equal(tr, [1.0, 1.0])
     t, y = dyn.expectations([[0,-1j],[1j,0]],real=True)
-    np.testing.assert_almost_equal(t, dyn.times)
-    np.testing.assert_almost_equal(y, [0, 0.2])
+    default_np.testing.assert_array_almost_equal(t, dyn.times)
+    default_np.testing.assert_array_almost_equal(y, [0, 0.2])
     t, lower = dyn.expectations([[0,0],[1,0]])
-    np.testing.assert_almost_equal(t, dyn.times)
-    np.testing.assert_almost_equal(lower, [0, -0.1j])
+    default_np.testing.assert_array_almost_equal(t, dyn.times)
+    default_np.testing.assert_array_almost_equal(lower, [0, -0.1j])
     with pytest.raises(AssertionError):
         dyn.expectations("bla")
 
@@ -97,14 +98,14 @@ def test_mean_field_dynamics():
                                      [states[1], states[1]]],
                                     fields)
     assert len(dyn_B) == len(times)
-    np.testing.assert_almost_equal(fields, dyn_B.fields)
+    default_np.testing.assert_array_almost_equal(fields, dyn_B.fields)
     dyn_B.add(times[-1] + other_time, [other_state, 2*other_state], other_field)
     assert isinstance(dyn_B.system_dynamics[0], Dynamics)
     sysB1, sysB2 = dyn_B.system_dynamics
-    np.testing.assert_almost_equal(sysB1.times, dyn_B.times)
-    np.testing.assert_almost_equal(sysB2.times, dyn_B.times)
+    default_np.testing.assert_array_almost_equal(sysB1.times, dyn_B.times)
+    default_np.testing.assert_array_almost_equal(sysB2.times, dyn_B.times)
     assert len(dyn_B) == len(sysB1)
-    np.testing.assert_almost_equal(sysB1.states[-1], 0.5 * sysB2.states[-1])
+    default_np.testing.assert_array_almost_equal(sysB1.states[-1], 0.5 * sysB2.states[-1])
     combined_fields = [fields[0], fields[1], other_field]
     # bad input
     with pytest.raises(AssertionError):
@@ -120,8 +121,8 @@ def test_mean_field_dynamics():
     assert t is None
     assert alpha is None
     t, alpha = dyn_B.field_expectations()
-    np.testing.assert_almost_equal(t, dyn_B.times)
-    np.testing.assert_almost_equal(alpha, dyn_B.fields)
-    np.testing.assert_almost_equal(alpha, combined_fields)
+    default_np.testing.assert_array_almost_equal(t, dyn_B.times)
+    default_np.testing.assert_array_almost_equal(alpha, dyn_B.fields)
+    default_np.testing.assert_array_almost_equal(alpha, combined_fields)
     with pytest.raises(TypeError):
         dyn_B.field_expectations("bla")
diff --git a/tests/coverage/gibbs_parameters_test.py b/tests/coverage/gibbs_parameters_test.py
index 36211b8..84ce4de 100644
--- a/tests/coverage/gibbs_parameters_test.py
+++ b/tests/coverage/gibbs_parameters_test.py
@@ -14,8 +14,6 @@
 """
 
 import pytest
-import numpy as np
-
 
 from oqupy import tempo
 
diff --git a/tests/coverage/gibbs_tempo_test.py b/tests/coverage/gibbs_tempo_test.py
index a530edf..2eb38d4 100644
--- a/tests/coverage/gibbs_tempo_test.py
+++ b/tests/coverage/gibbs_tempo_test.py
@@ -12,13 +12,15 @@
 """
 Tests for the Tempo module.
 """
+
 import pytest
-import numpy as np
+from numpy import ndarray
 
 import oqupy
 
 import oqupy.tempo as tempo
 
+from oqupy.backends.numerical_backend import default_np, np
 
 
 def test_gibbs_tempo():
@@ -64,31 +66,34 @@ def test_gibbs_tempo():
     state2 = state2 / state2.trace()
     assert (state1 == state2).all()
     assert state1.shape == (2, 2)
-    np.testing.assert_almost_equal(state1.trace(),1)
+    default_np.testing.assert_almost_equal(state1.trace(),1)
 
 # -- operators for testing purposes --
-def spin_z(n: int) -> np.ndarray:
+def spin_z(n: int) -> ndarray:
     sz = [n / 2]
     while len(sz) < n + 1:
         sz.append(sz[-1] - 1)
-    sz = np.diag(sz)
+    sz = np.diag(np.array(sz))
     return sz
 
-def spin_ladder_up(n: int) -> np.ndarray:
+def spin_ladder_up(n: int) -> ndarray:
     sp = np.zeros((n + 1, n + 1))
     for jj in range(len(sp) - 1):
-        sp[jj][jj + 1] = 0.5 * np.sqrt(
-            0.5 * n * (0.5 * n + 1)
-            - (0.5 * n - jj) * (0.5 * n - jj - 1))
+        sp = np.update(
+            array=sp,
+            indices=(jj, jj + 1),
+            values=0.5 * np.sqrt(0.5 * n * (0.5 * n + 1) - \
+                                 (0.5 * n - jj) * (0.5 * n - jj - 1))
+        )
     return sp
 
-def spin_ladder_down(n: int) -> np.ndarray:
+def spin_ladder_down(n: int) -> ndarray:
     return spin_ladder_up(n).T
 
-def spin_x(n: int) -> np.ndarray:
+def spin_x(n: int) -> ndarray:
     return spin_ladder_up(n) + spin_ladder_down(n)
 
-def spin_y(n: int) -> np.ndarray:
+def spin_y(n: int) -> ndarray:
     return 1j*(spin_ladder_up(n) - spin_ladder_down(n))
 
 def spin_operators(n: int) -> list:
diff --git a/tests/coverage/gradient_test.py b/tests/coverage/gradient_test.py
index 59dcff3..7679b36 100644
--- a/tests/coverage/gradient_test.py
+++ b/tests/coverage/gradient_test.py
@@ -14,8 +14,6 @@
 """
 
 import pytest
-import numpy as np
-from numpy import ndarray
 
 from typing import Dict
 
@@ -23,6 +21,8 @@
 
 from oqupy.gradient import state_gradient
 
+from oqupy.backends.numerical_backend import np
+
 def test_state_gradient():
     start_time=0
     num_steps=3
@@ -103,7 +103,7 @@ def discrete_h_sys(hx):
     assert callable(propagators)
     assert callable(propagator_derivatives)
     # Check the tot_derivs is a numpy array
-    assert isinstance(tot_derivs,ndarray)
+    assert isinstance(tot_derivs, np.ndarray)
 
 
 def test_compute_gradient_and_dynamics():
diff --git a/tests/coverage/node_array_test.py b/tests/coverage/node_array_test.py
index 71c4af7..05e28fa 100644
--- a/tests/coverage/node_array_test.py
+++ b/tests/coverage/node_array_test.py
@@ -17,19 +17,15 @@
 """
 import pytest
 
-import numpy as np
-import tensornetwork as tn
-
 import oqupy.backends.node_array as node_array
-
+from oqupy.backends.numerical_backend import np
 
 def test_node_array():
-    np.random.seed(0)
-    a = np.random.rand(3,4,9,5)
-    b = np.random.rand(5,5,9,6)
-    c = np.random.rand(6,5,9,2)
-    d = np.random.rand(2,3,9)
-    d2 = np.random.rand(2,3,9,3)
+    a = np.get_random_floats(0, (3,4,9,5))
+    b = np.get_random_floats(0, (5,5,9,6))
+    c = np.get_random_floats(0, (6,5,9,2))
+    d = np.get_random_floats(0, (2,3,9))
+    d2 = np.get_random_floats(0, (2,3,9,3))
     na = node_array.NodeArray([a,b,c,d], right=False, name="my node array")
     na2 = node_array.NodeArray([a,b,c,d2], name="my node array")
     str(na)
@@ -48,12 +44,12 @@ def test_node_array():
     dd2.get_verbose_string()
 
     #Apply matrix (left/right)
-    m = np.random.rand(3,7)
+    m = np.get_random_floats(0, (3,7))
     na.apply_matrix(m,left=True)
     na2.apply_matrix(m,left=False)
 
     # Apply vector (left/right)
-    v = np.random.rand(7)
+    v = np.get_random_floats(0, (7))
     na.apply_vector(v, left=True)
     na2.apply_vector(v, left=False)
 
@@ -61,12 +57,11 @@ def test_node_array():
     assert na.rank == None
 
 def test_node_array_join_split():
-    np.random.seed(0)
-    a = np.random.rand(3,4,9,5)
-    b = np.random.rand(5,5,9,6)
-    c = np.random.rand(6,5,9,2)
-    d = np.random.rand(2,3,9)
-    x = np.random.rand(2,9,3)
+    a = np.get_random_floats(0, (3,4,9,5))
+    b = np.get_random_floats(0, (5,5,9,6))
+    c = np.get_random_floats(0, (6,5,9,2))
+    d = np.get_random_floats(0, (2,3,9))
+    x = np.get_random_floats(0, (2,9,3))
 
     na0 = node_array.NodeArray([x], left=False, name="array 0")
     na1 = node_array.NodeArray([a,b], name="array 1")
@@ -80,11 +75,10 @@ def test_node_array_join_split():
         na_A, na_B = node_array.split(na_012, 5)
 
 def test_node_array_svd():
-    np.random.seed(0)
-    a = np.random.rand(3,4,9,5)
-    b = np.random.rand(5,5,9,6)
-    c = np.random.rand(6,5,9,2)
-    d = np.random.rand(2,3,9)
+    a = np.get_random_floats(0, (3,4,9,5))
+    b = np.get_random_floats(0, (5,5,9,6))
+    c = np.get_random_floats(0, (6,5,9,2))
+    d = np.get_random_floats(0, (2,3,9))
     na = node_array.NodeArray([a,b,c,d], right=False, name="my node array")
 
     na.svd_sweep(1, 3, max_singular_values=4)
@@ -98,20 +92,19 @@ def test_node_array_svd():
         na.svd_sweep(1, 5, max_singular_values=4)
 
 def test_node_array_contractions():
-    np.random.seed(0)
-    mps = node_array.NodeArray([np.random.rand(3,2,3),
-                                np.random.rand(3,3,3),
-                                np.random.rand(3,4,3)],
+    mps = node_array.NodeArray([np.get_random_floats(0, (3,2,3)),
+                                np.get_random_floats(0, (3,3,3)),
+                                np.get_random_floats(0, (3,4,3))],
                                name="MPS")
-    mpo1 = node_array.NodeArray([np.random.rand(2,5,4),
-                                 np.random.rand(4,3,5,4),
-                                 np.random.rand(4,4,5)],
+    mpo1 = node_array.NodeArray([np.get_random_floats(0, (2,5,4)),
+                                 np.get_random_floats(0, (4,3,5,4)),
+                                 np.get_random_floats(0, (4,4,5))],
                                 left=False,
                                 right=False,
                                 name="MPO1")
-    mpo2 = node_array.NodeArray([np.random.rand(5,5,3),
-                                 np.random.rand(3,5,5,3),
-                                 np.random.rand(3,5,5)],
+    mpo2 = node_array.NodeArray([np.get_random_floats(0, (5,5,3)),
+                                 np.get_random_floats(0, (3,5,5,3)),
+                                 np.get_random_floats(0, (3,5,5))],
                                 left=False,
                                 right=False,
                                 name="MPO2")
@@ -123,9 +116,9 @@ def test_node_array_contractions():
 
     # ### 2nd example
 
-    mps = node_array.NodeArray([np.random.rand(4,3),
-                                np.random.rand(3,4,3),
-                                np.random.rand(3,4)],
+    mps = node_array.NodeArray([np.get_random_floats(0, (4,3)),
+                                np.get_random_floats(0, (3,4,3)),
+                                np.get_random_floats(0, (3,4))],
                                left=False,
                                right=False,
                                name="MPS")
@@ -134,41 +127,41 @@ def test_node_array_contractions():
     mps3 = mps.copy()
     mps4 = mps.copy()
     mps5 = mps.copy()
-    mpsTwo = node_array.NodeArray([np.random.rand(4,3),
-                                 np.random.rand(3,4)],
+    mpsTwo = node_array.NodeArray([np.get_random_floats(0, (4,3)),
+                                 np.get_random_floats(0, (3,4))],
                                 left=False,
                                 right=False,
                                 name="MPS")
-    mpsL = node_array.NodeArray([np.random.rand(3,4,3),
-                                 np.random.rand(3,4)],
+    mpsL = node_array.NodeArray([np.get_random_floats(0, (3,4,3)),
+                                 np.get_random_floats(0, (3,4))],
                                 left=True,
                                 right=False,
                                 name="MPS")
-    mpsR = node_array.NodeArray([np.random.rand(4,3),
-                                 np.random.rand(3,4,3)],
+    mpsR = node_array.NodeArray([np.get_random_floats(0, (4,3)),
+                                 np.get_random_floats(0, (3,4,3))],
                                 left=False,
                                 right=True,
                                 name="MPS")
-    mpo1 = node_array.NodeArray([np.random.rand(3,4,4,3),
-                                 np.random.rand(3,4,4)],
+    mpo1 = node_array.NodeArray([np.get_random_floats(0, (3,4,4,3)),
+                                 np.get_random_floats(0, (3,4,4))],
                                 left=True,
                                 right=False,
                                 name="MPO1")
     mpo5 = mpo1.copy()
-    mpo2 = node_array.NodeArray([np.random.rand(4,4,3),
-                                 np.random.rand(3,4,4,3)],
+    mpo2 = node_array.NodeArray([np.get_random_floats(0, (4,4,3)),
+                                 np.get_random_floats(0, (3,4,4,3))],
                                 left=False,
                                 right=True,
                                 name="MPO2")
-    arr1 = node_array.NodeArray([np.random.rand(4,2,2,3),
-                                 np.random.rand(3,4,2,2,3),
-                                 np.random.rand(3,4,2,2,3)],
+    arr1 = node_array.NodeArray([np.get_random_floats(0, (4,2,2,3)),
+                                 np.get_random_floats(0, (3,4,2,2,3)),
+                                 np.get_random_floats(0, (3,4,2,2,3))],
                                 left=False,
                                 right=True,
                                 name="array1")
-    arr2 = node_array.NodeArray([np.random.rand(3,2,4,2,3),
-                                 np.random.rand(3,2,4,2,3),
-                                 np.random.rand(3,2,4,2)],
+    arr2 = node_array.NodeArray([np.get_random_floats(0, (3,2,4,2,3)),
+                                 np.get_random_floats(0, (3,2,4,2,3)),
+                                 np.get_random_floats(0, (3,2,4,2))],
                                 left=True,
                                 right=False,
                                 name="array2")
diff --git a/tests/coverage/operators_test.py b/tests/coverage/operators_test.py
index 8406f33..8666606 100644
--- a/tests/coverage/operators_test.py
+++ b/tests/coverage/operators_test.py
@@ -13,11 +13,10 @@
 Tests for the time_evovling_mpo.operators module.
 """
 
-import pytest
-import numpy as np
-
 from oqupy import operators
 
+from oqupy.backends.numerical_backend import default_np, np
+
 SIGMA = {"id":[[1, 0], [0, 1]],
          "x":[[0, 1], [1, 0]],
          "y":[[0, -1j], [1j, 0]],
@@ -39,32 +38,33 @@
 def test_identity():
     for n in {1,2,7}:
         result = operators.identity(n)
-        np.testing.assert_almost_equal(result,np.identity(n))
+        default_np.testing.assert_array_almost_equal(result,np.identity(n))
 
 def test_operators_sigma():
     for name in SIGMA:
         result = operators.sigma(name)
-        np.testing.assert_equal(result,SIGMA[name])
+        default_np.testing.assert_array_equal(result,SIGMA[name])
 
 def test_operators_spin_dm():
     for name in SPIN_DM:
         result = operators.spin_dm(name)
-        np.testing.assert_equal(result,SPIN_DM[name])
+        default_np.testing.assert_array_equal(result,SPIN_DM[name])
 
 def test_operators_create_destroy():
     for n in {1,2,7}:
         adag = operators.create(n)
         a = operators.destroy(n)
         result = adag@a
-        np.testing.assert_almost_equal(result,np.diag(range(n)))
+        default_np.testing.assert_array_almost_equal(result,np.diag(np.array(range(n), dtype=np.dtype_complex)))
 
 # -- testing super operators --------------------------------------------------
 
 N = 3
 
-a = np.random.rand(N,N) + 1j * np.random.rand(N,N)
-b = np.random.rand(N,N) + 1j * np.random.rand(N,N)
-x = np.random.rand(N,N) + 1j * np.random.rand(N,N)
+random_floats = np.get_random_floats(1234, (6, N, N))
+a = random_floats[0] + 1j * random_floats[1]
+b = random_floats[2] + 1j * random_floats[3]
+x = random_floats[4] + 1j * random_floats[5]
 
 a_dagger = a.conjugate().T
 b_dagger = a.conjugate().T
@@ -77,24 +77,24 @@ def test_operators_create_destroy():
 def test_commutator():
     sol = a@x - x@a
     res = operators.commutator(a)@x_vector
-    np.testing.assert_almost_equal(res.reshape(N,N),sol)
+    default_np.testing.assert_array_almost_equal(res.reshape(N,N),sol)
 
 def test_acommutator():
     sol = a@x + x@a
     res = operators.acommutator(a)@x_vector
-    np.testing.assert_almost_equal(res.reshape(N,N),sol)
+    default_np.testing.assert_array_almost_equal(res.reshape(N,N),sol)
 
 def test_left_super():
     sol = a@x
     res = operators.left_super(a)@x_vector
-    np.testing.assert_almost_equal(res.reshape(N,N),sol)
+    default_np.testing.assert_array_almost_equal(res.reshape(N,N),sol)
 
 def test_right_super():
     sol = x@a
     res = operators.right_super(a)@x_vector
-    np.testing.assert_almost_equal(res.reshape(N,N),sol)
+    default_np.testing.assert_array_almost_equal(res.reshape(N,N),sol)
 
 def test_left_right_super():
     sol = a@x@b
     res = operators.left_right_super(a,b)@x_vector
-    np.testing.assert_almost_equal(res.reshape(N,N),sol)
+    default_np.testing.assert_array_almost_equal(res.reshape(N,N),sol)
diff --git a/tests/coverage/pt_tebd_test.py b/tests/coverage/pt_tebd_test.py
index e05901c..5662b05 100644
--- a/tests/coverage/pt_tebd_test.py
+++ b/tests/coverage/pt_tebd_test.py
@@ -13,14 +13,13 @@
 Tests for the time_evovling_mpo.pt_tebd module.
 """
 
-import pytest
-
-import numpy as np
 import oqupy
 
+from oqupy.backends.numerical_backend import np
+
 up_dm = oqupy.operators.spin_dm("z+")
 system_chain = oqupy.SystemChain(hilbert_space_dimensions=[2,3])
-initial_augmented_mps = oqupy.AugmentedMPS([up_dm, np.diag([1,0,0])])
+initial_augmented_mps = oqupy.AugmentedMPS([up_dm, np.diag(np.array([1,0,0]))])
 pt_tebd_params = oqupy.PtTebdParameters(dt=0.2, order=2, epsrel=1.0e-4)
 
 def test_get_augmented_mps():
diff --git a/tests/coverage/pt_tempo_test.py b/tests/coverage/pt_tempo_test.py
index b9812f9..8119981 100644
--- a/tests/coverage/pt_tempo_test.py
+++ b/tests/coverage/pt_tempo_test.py
@@ -14,10 +14,11 @@
 """
 
 import pytest
-import numpy as np
 
 import oqupy as tempo
 
+from oqupy.backends.numerical_backend import np
+
 def test_pt_tempo():
     start_time = -0.3
     end_time = 0.84
diff --git a/tests/coverage/system_test.py b/tests/coverage/system_test.py
index 5c23376..fcfc361 100644
--- a/tests/coverage/system_test.py
+++ b/tests/coverage/system_test.py
@@ -14,12 +14,13 @@
 """
 
 import pytest
-import numpy as np
+from typing import Callable
 
 from oqupy.system import BaseSystem, System, TimeDependentSystem,\
         TimeDependentSystemWithField, MeanFieldSystem, ParameterizedSystem
 from oqupy import operators
-from typing import Callable
+
+from oqupy.backends.numerical_backend import default_np, np
 
 # -----------------------------------------------------------------------------
 # -- test-examples ------------------------------------------------------------
@@ -93,7 +94,7 @@ def test_base_system():
 def test_system_A():
     sys = System(hamiltonianA)
     liouvillian = sys.liouvillian()
-    np.testing.assert_almost_equal(liouvillian, liouvillianA)
+    default_np.testing.assert_array_almost_equal(liouvillian, liouvillianA)
 
 def test_system_B():
     sys = System(hamiltonianB,
@@ -109,15 +110,15 @@ def test_system_B():
     assert isinstance(sys.lindblad_operators[0], np.ndarray)
     assert sys.dimension == 2
     liouvillian = sys.liouvillian()
-    np.testing.assert_almost_equal(liouvillian, liouvillianB)
+    default_np.testing.assert_array_almost_equal(liouvillian, liouvillianB)
 
 def test_system_bad_input():
     with pytest.raises(AssertionError):
         System("bla")
     with pytest.raises(AssertionError):
-        System(np.random.rand(2,2,2))
+        System(np.get_random_floats(1, (2,2,2)))
     with pytest.raises(AssertionError):
-        System(np.random.rand(2,3))
+        System(np.get_random_floats(2, (2,3)))
     with pytest.raises(AssertionError):
         System(operators.sigma("z"),
                0.1,
@@ -144,7 +145,7 @@ def test_system_bad_input():
 def test_time_dependent_system_C():
     sys = TimeDependentSystem(hamiltonianC)
     liouvillian = sys.liouvillian(timeC)
-    np.testing.assert_almost_equal(liouvillian, liouvillianC)
+    default_np.testing.assert_array_almost_equal(liouvillian, liouvillianC)
 
 def test_time_dependent_system_D():
     sys = TimeDependentSystem(
@@ -157,23 +158,23 @@ def test_time_dependent_system_D():
         sys.liouvillian()
 
     str(sys)
-    assert isinstance(sys.hamiltonian, np.vectorize)
+    # assert isinstance(sys.hamiltonian, np.vectorize)
     assert isinstance(sys.gammas, list)
     assert sys.gammas[0](1.0) == gammasD[0](1.0)
     assert isinstance(sys.lindblad_operators, list)
-    assert isinstance(sys.lindblad_operators[0], np.vectorize)
+    # assert isinstance(sys.lindblad_operators[0], np.vectorize)
     assert sys.dimension == 2
     liouvillian = sys.liouvillian(timeD)
-    np.testing.assert_almost_equal(liouvillian, liouvillianD)
+    default_np.testing.assert_array_almost_equal(liouvillian, liouvillianD)
 
 
 def test_time_dependent_system_bad_input():
     with pytest.raises(AssertionError):
         TimeDependentSystem("bla")
     with pytest.raises(AssertionError):
-        TimeDependentSystem(lambda t: t*np.random.rand(2,2,2))
+        TimeDependentSystem(lambda t: t*np.get_random_floats(1, (2,2,2)))
     with pytest.raises(AssertionError):
-        TimeDependentSystem(lambda t: t*np.random.rand(2,3))
+        TimeDependentSystem(lambda t: t*np.get_random_floats(1, (2,3)))
     with pytest.raises(AssertionError):
         TimeDependentSystem(lambda t: t*operators.sigma("z"),
                0.1,
@@ -203,7 +204,7 @@ def test_time_dependent_system_with_field():
             name="with field")
     liouvillian = sysField.liouvillian(timeField,
             timeField+dtField, fieldField, derivativeField)
-    np.testing.assert_almost_equal(liouvillian, liouvillianField)
+    default_np.testing.assert_array_almost_equal(liouvillian, liouvillianField)
     # Liouvillian calling
     with pytest.raises(TypeError):
         sysField.liouvillian()
diff --git a/tests/coverage/tempo_parameters_test.py b/tests/coverage/tempo_parameters_test.py
index 7e91037..88aa8fd 100644
--- a/tests/coverage/tempo_parameters_test.py
+++ b/tests/coverage/tempo_parameters_test.py
@@ -14,10 +14,11 @@
 """
 
 import pytest
-import numpy as np
 
 import oqupy as tempo
 
+from oqupy.backends.numerical_backend import np
+
 def test_tempo_parameters():
     tempo_param = tempo.TempoParameters(
         0.1, 1.0e-5, None, None, None, None, 2.0e-5, "rough", "bla")
diff --git a/tests/coverage/tempo_test.py b/tests/coverage/tempo_test.py
index 95cec34..72420d9 100644
--- a/tests/coverage/tempo_test.py
+++ b/tests/coverage/tempo_test.py
@@ -14,10 +14,11 @@
 """
 
 import pytest
-import numpy as np
 
 import oqupy
 
+from oqupy.backends.numerical_backend import np
+
 def test_tempo():
     start_time = -0.3
     end_time1 = 0.4
diff --git a/tests/physics/__init__.py b/tests/physics/__init__.py
index e69de29..9ae2946 100644
--- a/tests/physics/__init__.py
+++ b/tests/physics/__init__.py
@@ -0,0 +1,2 @@
+# from oqupy.backends import enable_jax_features
+# enable_jax_features()
diff --git a/tests/physics/bath_dynamics_test.py b/tests/physics/bath_dynamics_test.py
index 2d0372a..8f5ef7e 100644
--- a/tests/physics/bath_dynamics_test.py
+++ b/tests/physics/bath_dynamics_test.py
@@ -12,12 +12,13 @@
 """
 Tests for the oqupy.bath_dynamics module.
 """
-import numpy as np
 
 from oqupy.bath_dynamics import TwoTimeBathCorrelations
 from oqupy.bath_correlations import PowerLawSD
-from oqupy import TempoParameters, System, Bath, PtTempo
-from oqupy.config import NpDtype
+from oqupy.config import NumPyDtypeComplex
+from oqupy import Bath, PtTempo, System, TempoParameters
+
+from oqupy.backends.numerical_backend import default_np, np
 
 def exact_correlation(t_1, t_2, w_1, w_2, dagg,
                g_1, g_2, temp):
@@ -51,7 +52,7 @@ def test():
     final_time = 2
     epsrel = 1e-7
     system_correlations = np.ones((int(final_time/dt),                  #system correlations
-                              int(final_time/dt)), dtype = NpDtype)     #are known exactly.
+                              int(final_time/dt)), dtype=NumPyDtypeComplex)     #are known exactly.
     correlations = PowerLawSD(alpha=alpha,
                               zeta=1.0,
                               cutoff=wc,
@@ -70,12 +71,12 @@ def test():
     #Test properties
     assert corr.system == system
     assert corr.bath == bath
-    np.testing.assert_equal(corr.initial_state, initial_state)
+    default_np.testing.assert_array_equal(corr.initial_state, initial_state)
 
     tlist, occ_0 = corr.occupation(0)
 
-    np.testing.assert_equal(occ_0,np.ones(len(tlist),
-                                        dtype=NpDtype) * (np.nan + 1.0j*np.nan))
+    default_np.testing.assert_array_equal(occ_0,np.ones(len(tlist),
+                                        dtype=NumPyDtypeComplex) * (np.nan + 1.0j*np.nan))
 
     w0 = 1
     coup = correlations.spectral_density(w0)
@@ -107,10 +108,9 @@ def test():
                                                  temperature))
         int_phase = np.exp(-1j * ((2 * dagg[0] - 1) * w02 * tlist[sel:] + \
                                   (2 * dagg[1] - 1) * w0 * tlist[sel]))
-        assert np.allclose(corrs,exact_corrs)
-        assert np.allclose(int_corrs,exact_corrs * int_phase)
+        default_np.testing.assert_array_almost_equal(corrs, exact_corrs)
+        default_np.testing.assert_array_almost_equal(int_corrs, np.array(exact_corrs) * int_phase)
     corr2 = TwoTimeBathCorrelations(system, bath, pt,
                                     initial_state=initial_state)
     corr2.generate_system_correlations(1)
-    assert np.allclose(corr2._system_correlations[0,:],
-                       np.ones(int(len(tlist)/2)))
+    assert np.allclose(corr2._system_correlations[0,:], np.ones(int(len(tlist)/2)))
diff --git a/tests/physics/degeneracy_large_test.py b/tests/physics/degeneracy_large_test.py
index 4354d11..ee3cafc 100644
--- a/tests/physics/degeneracy_large_test.py
+++ b/tests/physics/degeneracy_large_test.py
@@ -14,10 +14,10 @@
 """
 
 import pytest
-import numpy as np
 
 import oqupy
-from oqupy import process_tensor
+
+from oqupy.backends.numerical_backend import default_np, np
 
 
 # -----------------------------------------------------------------------------
@@ -40,7 +40,7 @@
 h_sys_C = (j_value/2) * s_z @ s_z + h * s_x
 
 # Ohmic spectral density with exponential cutoff
-coupling_operator_C = np.diag([1,1,2])
+coupling_operator_C = np.diag(np.array([1,1,2], dtype=np.dtype_complex))
 alpha_C = 0.3
 cutoff_C = 5.0
 temperature_C = 0.2
@@ -83,6 +83,5 @@ def test_large_degeneracy_compare():
     tempo_non_unique.compute(end_time=t_end_C)
     dyn_unique = tempo_unique.get_dynamics()
     dyn_non_unique = tempo_non_unique.get_dynamics()
-    np.testing.assert_almost_equal(dyn_unique.states, dyn_non_unique.states,
-                                   decimal=4)
+    default_np.testing.assert_array_almost_equal(dyn_unique.states, dyn_non_unique.states, decimal=4)
 # -----------------------------------------------------------------------------
diff --git a/tests/physics/degeneracy_mean_field_test.py b/tests/physics/degeneracy_mean_field_test.py
index 782a160..0ebb6ef 100644
--- a/tests/physics/degeneracy_mean_field_test.py
+++ b/tests/physics/degeneracy_mean_field_test.py
@@ -13,11 +13,10 @@
 Tests degeneracy checking with mean field TEMPO.
 """
 
-import pytest
-import numpy as np
-
 import oqupy
 
+from oqupy.backends.numerical_backend import default_np, np
+
 # -----------------------------------------------------------------------------
 # -- Test A: Spin boson model -------------------------------------------------
 
@@ -226,7 +225,7 @@ def test_degeneracy_exact():
         unique=True)
     tempo_A.compute(end_time=t_end_A)
     dyn_A = tempo_A.get_dynamics()
-    np.testing.assert_almost_equal(dyn_A.states[-1], rho_A, decimal=4)
+    default_np.testing.assert_array_almost_equal(dyn_A.states[-1], rho_A, decimal=4)
 
 def test_degeneracy_compare():
     tempo_params_B = oqupy.TempoParameters(
@@ -251,8 +250,7 @@ def test_degeneracy_compare():
     tempo_non_unique.compute(end_time=t_end_B)
     dyn_unique = tempo_unique.get_dynamics()
     dyn_non_unique = tempo_non_unique.get_dynamics()
-    np.testing.assert_almost_equal(dyn_unique.states, dyn_non_unique.states,
-                                   decimal=4)
+    default_np.testing.assert_array_almost_equal(dyn_unique.states, dyn_non_unique.states, decimal=4)
 
 def test_degeneracy_1d():
     tempo_params_C = oqupy.TempoParameters(
@@ -277,8 +275,7 @@ def test_degeneracy_1d():
     tempo_non_unique.compute(end_time=t_end_C)
     dyn_unique = tempo_unique.get_dynamics()
     dyn_non_unique = tempo_non_unique.get_dynamics()
-    np.testing.assert_almost_equal(dyn_unique.states, dyn_non_unique.states,
-                                   decimal=4)
+    default_np.testing.assert_array_almost_equal(dyn_unique.states, dyn_non_unique.states, decimal=4)
 
 def test_mean_field_degeneracy_compare():
     tempo_params_D = oqupy.TempoParameters(
@@ -307,10 +304,8 @@ def test_mean_field_degeneracy_compare():
     _, field_unique = dyn1.field_expectations()
     dyn_non_unique = dyn2.system_dynamics[0]
     _, field_non_unique = dyn2.field_expectations()
-    np.testing.assert_almost_equal(dyn_unique.states, dyn_non_unique.states,
-                                   decimal=4)
-    np.testing.assert_almost_equal(field_unique, field_non_unique,
-                                   decimal=4)
+    default_np.testing.assert_array_almost_equal(dyn_unique.states, dyn_non_unique.states, decimal=4)
+    default_np.testing.assert_array_almost_equal(field_unique, field_non_unique, decimal=4)
     
 def test_pt_tempo_degeneracy_compare():
     tempo_params_E = oqupy.TempoParameters(
@@ -336,6 +331,5 @@ def test_pt_tempo_degeneracy_compare():
     dyn_non_unique = oqupy.compute_dynamics(system=system_E, 
                            initial_state=initial_state_E,
                            process_tensor=pt_non_unique)
-    np.testing.assert_almost_equal(dyn_unique.states, dyn_non_unique.states,
-                                   decimal=4)
+    default_np.testing.assert_array_almost_equal(dyn_unique.states, dyn_non_unique.states, decimal=4)
 # -----------------------------------------------------------------------------
diff --git a/tests/physics/gibbs_tempo_test.py b/tests/physics/gibbs_tempo_test.py
index ec4c589..a6ed7d9 100644
--- a/tests/physics/gibbs_tempo_test.py
+++ b/tests/physics/gibbs_tempo_test.py
@@ -14,10 +14,11 @@
 """
 
 import pytest
-import numpy as np
 
 import oqupy
 
+from oqupy.backends.numerical_backend import default_np, np
+
 # -----------------------------------------------------------------------------
 
 h_sys = np.array([[0.5, 0.0], [0, -0.5]])
@@ -98,14 +99,14 @@ def test_gibbs_backend_A():
         gibbs_params)
     tempo_A.compute()
     gibbs_A = tempo_A.get_state()
-    np.testing.assert_almost_equal(gibbs_A, indep_sbm(alpha_A), decimal=8)
+    default_np.testing.assert_array_almost_equal(gibbs_A, indep_sbm(alpha_A), decimal=8)
 
 def test_gibbs_backend_B():
     gibbs_B = oqupy.gibbs_tempo_compute(
         system_B,
         bath_B,
         gibbs_params)
-    np.testing.assert_almost_equal(gibbs_B, indep_sbm(alpha_B), decimal=8)
+    default_np.testing.assert_array_almost_equal(gibbs_B, indep_sbm(alpha_B), decimal=8)
 
 @pytest.mark.skip(reason="See GitHub Issue#???")
 def test_gibbs_backend_C():
@@ -113,6 +114,6 @@ def test_gibbs_backend_C():
         system_C,
         bath_C,
         gibbs_params)
-    np.testing.assert_almost_equal(gibbs_C, transf(indep_sbm(alpha_C)), decimal=8)
+    default_np.testing.assert_array_almost_equal(gibbs_C, transf(indep_sbm(alpha_C)), decimal=8)
 
 # -----------------------------------------------------------------------------
diff --git a/tests/physics/gradient_functional_target_test.py b/tests/physics/gradient_functional_target_test.py
index 50058a8..cfed6e5 100644
--- a/tests/physics/gradient_functional_target_test.py
+++ b/tests/physics/gradient_functional_target_test.py
@@ -13,12 +13,9 @@
 Tests for the compute_gradient_and_dynamics and chain_rule functions.
 """
 
-import pytest
-import numpy as np
-
 import oqupy
-from oqupy import process_tensor
 
+from oqupy.backends.numerical_backend import default_np, np
 
 # -----------------------------------------------------------------------------
 # -- Test K: Spin-Boson model gradient (time-dep. sys., functional target state)
@@ -104,7 +101,7 @@ def test_tempo_gradient_backend_K():
         progress_type='silent'
     )
 
-    np.testing.assert_almost_equal(dyn.states[-1], rho_K, decimal=4)
+    default_np.testing.assert_array_almost_equal(dyn.states[-1], rho_K, decimal=4)
 
     get_props = system_K.get_propagators(dt,x0)
     get_prop_derivatives = system_K.get_propagator_derivatives(dt,x0)
@@ -115,4 +112,4 @@ def test_tempo_gradient_backend_K():
                                             num_steps=num_steps,
                                             num_parameters=1)
     
-    np.testing.assert_almost_equal(grad_params.real,grad_params_K,decimal=4)
\ No newline at end of file
+    default_np.testing.assert_array_almost_equal(grad_params.real,grad_params_K,decimal=4)
\ No newline at end of file
diff --git a/tests/physics/gradient_target_state_test.py b/tests/physics/gradient_target_state_test.py
index 54c6aeb..af339d4 100644
--- a/tests/physics/gradient_target_state_test.py
+++ b/tests/physics/gradient_target_state_test.py
@@ -13,11 +13,9 @@
 Tests for the compute_gradient_and_dynamics and chain_rule functions.
 """
 
-import pytest
-import numpy as np
-
 import oqupy
 
+from oqupy.backends.numerical_backend import default_np, np
 
 # -----------------------------------------------------------------------------
 # -- Test J: Spin boson model gradient -------------------------------------------------
@@ -99,7 +97,7 @@ def test_tempo_gradient_backend_J():
                                                     target_derivative=target_derivative_J.T
                                                     )
 
-    np.testing.assert_almost_equal(dyn.states[-1], rho_J, decimal=4)
+    default_np.testing.assert_array_almost_equal(dyn.states[-1], rho_J, decimal=4)
 
     get_props = system_J.get_propagators(dt,parameters=x0)
     get_prop_derivatives = system_J.get_propagator_derivatives(dt=dt,parameters=x0)
@@ -110,4 +108,4 @@ def test_tempo_gradient_backend_J():
                                             num_steps=num_steps,
                                             num_parameters=1)
     
-    np.testing.assert_almost_equal(grad_params.real,grad_params_J,decimal=4)
+    default_np.testing.assert_array_almost_equal(grad_params.real,grad_params_J,decimal=4)
diff --git a/tests/physics/mean_field_tempo_multi_test.py b/tests/physics/mean_field_tempo_multi_test.py
index 1bd6df5..53cc297 100644
--- a/tests/physics/mean_field_tempo_multi_test.py
+++ b/tests/physics/mean_field_tempo_multi_test.py
@@ -1,10 +1,11 @@
-import numpy as np
-import oqupy
-
 """
 Test dynamics with multiple identical mean-field systems are consistent.
 """
 
+import oqupy
+
+from oqupy.backends.numerical_backend import np
+
 number_of_systems = 3
 dt = 0.001
 t_end_I = 30 * dt
diff --git a/tests/physics/mean_field_tempo_test.py b/tests/physics/mean_field_tempo_test.py
index c7cd068..7332d18 100644
--- a/tests/physics/mean_field_tempo_test.py
+++ b/tests/physics/mean_field_tempo_test.py
@@ -14,11 +14,12 @@
 """
 Tests for MeanFieldTempo.
 """
-import numpy as np
 
 import oqupy
 from oqupy import operators
 
+from oqupy.backends.numerical_backend import default_np, np
+
 # -- Test G: Dicke model --
 
 initial_state_G = np.array([[0.5, 0.5], [0.5, 0.5]])
@@ -67,8 +68,8 @@ def test_tempo_backend_G():
         start_time=0.0)
     tempo_G.compute(end_time=t_end_G)
     dyn_G = tempo_G.get_dynamics()
-    np.testing.assert_almost_equal(dyn_G.fields[-1], field_G, decimal=4)
-    np.testing.assert_almost_equal(
+    default_np.testing.assert_array_almost_equal(dyn_G.fields[-1], field_G, decimal=4)
+    default_np.testing.assert_array_almost_equal(
             dyn_G.system_dynamics[0].states[-1], rho_G, decimal=4)
 
 def test_tensor_network_pt_tempo_backend_A():
@@ -83,8 +84,8 @@ def test_tensor_network_pt_tempo_backend_A():
         initial_field=initial_field_G,
         process_tensor_list=[pt],
         initial_state_list=[initial_state_G])
-    np.testing.assert_almost_equal(dyn.fields[-1], field_G, decimal=4)
-    np.testing.assert_almost_equal(
+    default_np.testing.assert_array_almost_equal(dyn.fields[-1], field_G, decimal=4)
+    default_np.testing.assert_array_almost_equal(
             dyn.system_dynamics[0].states[-1], rho_G, decimal=4)
 
 # -----------------------------------------------------------------------------
diff --git a/tests/physics/multi_environments_test.py b/tests/physics/multi_environments_test.py
index 0e842d3..4b642f4 100644
--- a/tests/physics/multi_environments_test.py
+++ b/tests/physics/multi_environments_test.py
@@ -13,10 +13,9 @@
 Test for multiple environments.
 """
 
-import numpy as np
-
 import oqupy
 
+from oqupy.backends.numerical_backend import default_np, np
 
 # -- prepare process tensors -------------------------------------------------
 
@@ -65,7 +64,7 @@ def test_multi_env_dynamics():
     dyns2 = oqupy.compute_dynamics(system,
                                    process_tensor=[pt2],
                                    initial_state=initial_state)
-    np.testing.assert_almost_equal(dyns.states,dyns2.states,decimal=5)
+    default_np.testing.assert_array_almost_equal(dyns.states,dyns2.states,decimal=5)
 
 
 def test_multi_env_gradient():
@@ -89,4 +88,4 @@ def test_multi_env_gradient():
                                 dt=dt,
                                 parameters=parameter_list)
     
-    np.testing.assert_almost_equal(gradsA[0],gradsB[0],decimal=5)
+    default_np.testing.assert_array_almost_equal(gradsA[0],gradsB[0],decimal=5)
diff --git a/tests/physics/multi_time_correlations_test.py b/tests/physics/multi_time_correlations_test.py
index 0e42410..362df6a 100644
--- a/tests/physics/multi_time_correlations_test.py
+++ b/tests/physics/multi_time_correlations_test.py
@@ -15,11 +15,9 @@
 import sys
 sys.path.insert(0,'.')
 
-import pytest
-import numpy as np
-
 import oqupy
 
+from oqupy.backends.numerical_backend import np
 
 # -----------------------------------------------------------------------------
 # -- Test: Multi-time correlations for independent boson model  ---------------
@@ -67,7 +65,7 @@
 g_t=[]
 t_corr0=np.arange(0.,end_time+dt,dt)
 for i in range(len(t_corr0)):
-    g=correlations.correlation_2d_integral(delta=t_corr0[i],
+    g=correlations.correlation_2d_integral(delta=float(t_corr0[i]),
                                            time_1=0.,shape="upper-triangle")
     g_t.append(g)
 
diff --git a/tests/physics/pt_tebd_lindblad_test.py b/tests/physics/pt_tebd_lindblad_test.py
index 09175b3..ae98ead 100644
--- a/tests/physics/pt_tebd_lindblad_test.py
+++ b/tests/physics/pt_tebd_lindblad_test.py
@@ -16,11 +16,10 @@
 import sys
 sys.path.insert(0,'.')
 
-import pytest
-import numpy as np
-
 import oqupy
 
+from oqupy.backends.numerical_backend import default_np, np
+
 # -----------------------------------------------------------------------------
 # -- Test F: Test Lindblad dissipation for PT-TEBD  ---------------------------
 
@@ -83,15 +82,15 @@ def test_pt_tebd_site_dissipation_H1():
 
     r = pt_tebd.compute(num_steps, progress_type="silent")
 
-    np.testing.assert_almost_equal(
+    default_np.testing.assert_array_almost_equal(
         r['dynamics'][0].states[-1],
         [[np.exp(-1),0],[0,1-np.exp(-1)]],
         decimal=4)
-    np.testing.assert_almost_equal(
+    default_np.testing.assert_array_almost_equal(
         r['dynamics'][1].states[-1],
         [[0,0],[0,1]],
         decimal=4)
-    np.testing.assert_almost_equal(
+    default_np.testing.assert_array_almost_equal(
         r['dynamics'][2].states[-1],
         [[1-np.exp(-1),0],[0,np.exp(-1)]],
         decimal=4)
diff --git a/tests/physics/pt_tebd_test.py b/tests/physics/pt_tebd_test.py
index 1dc6ffa..deeea43 100644
--- a/tests/physics/pt_tebd_test.py
+++ b/tests/physics/pt_tebd_test.py
@@ -15,11 +15,10 @@
 import sys
 sys.path.insert(0,'.')
 
-import pytest
-import numpy as np
-
 import oqupy
 
+from oqupy.backends.numerical_backend import default_np, np
+
 # -----------------------------------------------------------------------------
 # -- Test F: XYZ Spin chain without and with bath  ----------------------------
 
@@ -121,7 +120,7 @@ def test_pt_tebd_backend_F1():
     for site in dynamics_sites:
         pt_tebd_rho = r['dynamics'][site].states[-1]
         correct_rho = correct_F1_rhos[site]
-        np.testing.assert_almost_equal(pt_tebd_rho, correct_rho, decimal=4)
+        default_np.testing.assert_array_almost_equal(pt_tebd_rho, correct_rho, decimal=4)
 
 # -----------------------------------------------------------------------------
 
@@ -139,6 +138,6 @@ def test_pt_tebd_backend_F2():
     for site in dynamics_sites:
         pt_tebd_rho = r['dynamics'][site].states[-1]
         correct_rho = correct_F2_rhos[site]
-        np.testing.assert_almost_equal(pt_tebd_rho, correct_rho, decimal=4)
+        default_np.testing.assert_array_almost_equal(pt_tebd_rho, correct_rho, decimal=4)
 
 # -----------------------------------------------------------------------------
diff --git a/tests/physics/tempo_lindblad_test.py b/tests/physics/tempo_lindblad_test.py
index 1a9027c..98813a7 100644
--- a/tests/physics/tempo_lindblad_test.py
+++ b/tests/physics/tempo_lindblad_test.py
@@ -13,11 +13,10 @@
 Tests for TEMPO and PT-TEMPO with finite temperature and Lindblad dissipators.
 """
 
-import pytest
-import numpy as np
-
 import oqupy
 
+from oqupy.backends.numerical_backend import default_np, np
+
 # -----------------------------------------------------------------------------
 # -- Test A: Spin boson model -------------------------------------------------
 
@@ -75,7 +74,7 @@ def test_tensor_network_tempo_backend_D():
     tempo_D.compute(end_time=1.0)
     dyn_D = tempo_D.get_dynamics()
     print(dyn_D.states[-1])
-    np.testing.assert_almost_equal(dyn_D.states[-1], rho_D, decimal=4)
+    default_np.testing.assert_array_almost_equal(dyn_D.states[-1], rho_D, decimal=4)
 
 
 def test_tensor_network_pt_tempo_backend_D():
@@ -94,5 +93,5 @@ def test_tensor_network_pt_tempo_backend_D():
         process_tensor=pt,
         initial_state=initial_state_D)
     print(dyn.states[-1])
-    np.testing.assert_almost_equal(dyn.states[-1], rho_D, decimal=4)
+    default_np.testing.assert_array_almost_equal(dyn.states[-1], rho_D, decimal=4)
 # -----------------------------------------------------------------------------
diff --git a/tests/physics/tempo_non_diagonal_test.py b/tests/physics/tempo_non_diagonal_test.py
index 85a0352..b303255 100644
--- a/tests/physics/tempo_non_diagonal_test.py
+++ b/tests/physics/tempo_non_diagonal_test.py
@@ -13,11 +13,10 @@
 Tests for TEMPO and PT-TEMPO with non-diagonal coupling operators.
 """
 
-import pytest
-import numpy as np
-
 import oqupy
 
+from oqupy.backends.numerical_backend import np
+
 # -----------------------------------------------------------------------------
 # -- Test B: non-diagonal coupling --------------------------------------------
 
diff --git a/tests/physics/tempo_qutrit_test.py b/tests/physics/tempo_qutrit_test.py
index 4492fa4..fca410d 100644
--- a/tests/physics/tempo_qutrit_test.py
+++ b/tests/physics/tempo_qutrit_test.py
@@ -13,11 +13,10 @@
 Tests for TEMPO and PT-TEMPO with a qudit.
 """
 
-import pytest
-import numpy as np
-
 import oqupy
 
+from oqupy.backends.numerical_backend import default_np, np
+
 # -----------------------------------------------------------------------------
 # -- Test C: Spin-1 boson model -------------------------------------------------
 
@@ -92,7 +91,7 @@ def test_tensor_network_tempo_backend_C():
     tempo_C.compute(end_time=1.0)
     dyn_C = tempo_C.get_dynamics()
     assert dyn_C.times[-1] == 1.0
-    np.testing.assert_almost_equal(dyn_C.states[-1], rho_C, decimal=4)
+    default_np.testing.assert_array_almost_equal(dyn_C.states[-1], rho_C, decimal=4)
 
 
 def test_tensor_network_pt_tempo_backend_C():
@@ -112,6 +111,6 @@ def test_tensor_network_pt_tempo_backend_C():
         process_tensor=pt,
         initial_state=initial_state_C)
     assert dyn.times[-1] == 1.0
-    np.testing.assert_almost_equal(dyn.states[-1], rho_C, decimal=4)
+    default_np.testing.assert_array_almost_equal(dyn.states[-1], rho_C, decimal=4)
 
 # -----------------------------------------------------------------------------
diff --git a/tests/physics/tempo_spin_boson_test.py b/tests/physics/tempo_spin_boson_test.py
index 362d232..e76d0dd 100644
--- a/tests/physics/tempo_spin_boson_test.py
+++ b/tests/physics/tempo_spin_boson_test.py
@@ -13,11 +13,9 @@
 Tests for TEMPO and PT-TEMPO.
 """
 
-import pytest
-import numpy as np
-
 import oqupy
 
+from oqupy.backends.numerical_backend import default_np, np
 
 # -----------------------------------------------------------------------------
 # -- Test A: Spin boson model -------------------------------------------------
@@ -74,7 +72,7 @@ def test_tempo_backend_A():
         start_time=0.0)
     tempo_A.compute(end_time=t_end_A)
     dyn_A = tempo_A.get_dynamics()
-    np.testing.assert_almost_equal(dyn_A.states[-1], rho_A, decimal=4)
+    default_np.testing.assert_array_almost_equal(dyn_A.states[-1], rho_A, decimal=4)
 
 def test_tensor_network_pt_tempo_backend_A():
     tempo_params_A = oqupy.TempoParameters(
@@ -91,6 +89,6 @@ def test_tensor_network_pt_tempo_backend_A():
         system=system_A,
         process_tensor=pt,
         initial_state=initial_state_A)
-    np.testing.assert_almost_equal(dyn.states[-1], rho_A, decimal=4)
+    default_np.testing.assert_array_almost_equal(dyn.states[-1], rho_A, decimal=4)
 
 # -----------------------------------------------------------------------------
diff --git a/tests/physics/tempo_superohmic_test.py b/tests/physics/tempo_superohmic_test.py
index d883dc3..e91808f 100644
--- a/tests/physics/tempo_superohmic_test.py
+++ b/tests/physics/tempo_superohmic_test.py
@@ -13,11 +13,10 @@
 Tests TEMPO and PT-TEMPO with superohmic spectral density.
 """
 
-import pytest
-import numpy as np
-
 import oqupy
 
+from oqupy.backends.numerical_backend import default_np, np
+
 # -----------------------------------------------------------------------------
 # -- Test E: Superohmic independent spin boson model --------------------------
 
@@ -77,7 +76,7 @@ def test_tensor_network_tempo_backend_E():
         start_time=t_start_E)
     tempo_E.compute(end_time=t_end_E)
     dyn_E = tempo_E.get_dynamics()
-    np.testing.assert_almost_equal(dyn_E.states[-1], rho_E, decimal=4)
+    default_np.testing.assert_array_almost_equal(dyn_E.states[-1], rho_E, decimal=4)
 
 
 def test_tensor_network_pt_tempo_backend_E():
@@ -96,6 +95,6 @@ def test_tensor_network_pt_tempo_backend_E():
         system=system_E,
         process_tensor=pt,
         initial_state=initial_state_E)
-    np.testing.assert_almost_equal(dyn.states[-1], rho_E, decimal=4)
+    default_np.testing.assert_array_almost_equal(dyn.states[-1], rho_E, decimal=4)
 
 # -----------------------------------------------------------------------------