Option to manually choose between computation backends

[cjcarver]

I think it would be nice to not have to completely change my virtual environment every time I want to test simphony with or without JAX installed. So I was thinking we could just have a environment variable available for those who have JAX installed to optionally use the numpy environment. I was thinking something like this:

import os

backend = os.environ.get("SIMPHONY_BACKEND", "numpy")

JAX_AVAILABLE = False
if backend == "jax":
    try:
        import jax
        import jax.numpy as jnp
        JAX_AVAILABLE = True
    except ImportError:
        print("JAX not found, falling back to NumPy backend")
        backend = "numpy"
        
if backend == "numpy":
    import numpy as jnp
    from simphony.utils import jax 

[sequoiap]

I'm wondering about the usefulness of this. It's really only useful when testing for developers, correct? I cannot imagine a real-world use case where someone would be switching back and forth between the computational backends. And, even if they were, that's just as simple as having two terminals open, one of each environment, no? For some reason, I just hate complicating the import statement even more, since it has to be changed in every file and has to remain in sync between them...

[Andeloth]

I agree that the usefulness is limited, but I am actually more concerned about the import statement having to be pasted at the beginning of every file. I know imports are usually not modularized, but in this case would it make sense to have a function that dynamically imports the correct backend? Something like this https://stackoverflow.com/questions/301134/how-can-i-import-a-module-dynamically-given-its-name-as-string

And then it would be very simple to add the extra backend option as well if still desired.

Or maybe it should just be left as is and developer time should be better used elsewhere /shrug.