Switching and Verifying Backends¶

wdm-transform uses a small backend abstraction so the same code can run with NumPy, JAX, or CuPy. The default is NumPy, but you can select another backend explicitly or by setting the WDM_BACKEND environment variable.

This notebook shows the three practical checks: how to choose the backend, how to construct objects on that backend, and how to confirm which backend a created object is actually using.

In [1]:

import os

import wdm_transform as wt

backend = wt.get_backend()
print(f"Default backend: {backend.name}")

series = wt.TimeSeries([0.0, 1.0, 0.0, -1.0, 0.0, 1.0, 0.0, -1.0], dt=0.5, backend=backend)
print(f"Series backend: {series.backend.name}")
print(f"Underlying array type: {type(series.data).__module__}")

assert series.backend.name == backend.name
assert series.backend is backend

import os import wdm_transform as wt backend = wt.get_backend() print(f"Default backend: {backend.name}") series = wt.TimeSeries([0.0, 1.0, 0.0, -1.0, 0.0, 1.0, 0.0, -1.0], dt=0.5, backend=backend) print(f"Series backend: {series.backend.name}") print(f"Underlying array type: {type(series.data).__module__}") assert series.backend.name == backend.name assert series.backend is backend

Default backend: numpy
Series backend: numpy
Underlying array type: numpy

Change the active backend¶

You can resolve a backend by name with wt.get_backend("jax") or wt.get_backend("cupy"). If the corresponding dependency is not installed, the call raises an informative error.

In [2]:

for name in ["numpy", "jax", "cupy"]:
    try:
        resolved = wt.get_backend(name)
        print(f"{name!r} is available: {resolved.name}")
    except (ImportError, ValueError):
        print(f"{name!r} is not available in this environment")

for name in ["numpy", "jax", "cupy"]: try: resolved = wt.get_backend(name) print(f"{name!r} is available: {resolved.name}") except (ImportError, ValueError): print(f"{name!r} is not available in this environment")

'numpy' is available: numpy

'jax' is available: jax
'cupy' is not available in this environment

Set an environment default¶

If you prefer configuration via environment variables, set WDM_BACKEND before creating the first object or before calling wt.get_backend() with no arguments.

In [3]:

os.environ["WDM_BACKEND"] = "numpy"
resolved_default = wt.get_backend()
print(f"Environment default: {resolved_default.name}")
assert resolved_default.name == "numpy"

# A backend can also be given directly to each object.
series_from_env = wt.TimeSeries([1.0, 2.0, 3.0], dt=0.1, backend=resolved_default)
print(f"Object confirms it is on: {series_from_env.backend.name}")
assert series_from_env.backend is resolved_default

os.environ["WDM_BACKEND"] = "numpy" resolved_default = wt.get_backend() print(f"Environment default: {resolved_default.name}") assert resolved_default.name == "numpy" # A backend can also be given directly to each object. series_from_env = wt.TimeSeries([1.0, 2.0, 3.0], dt=0.1, backend=resolved_default) print(f"Object confirms it is on: {series_from_env.backend.name}") assert series_from_env.backend is resolved_default

Environment default: numpy
Object confirms it is on: numpy

Check backend state in practice¶

A good defensive pattern is to resolve the current backend once and then confirm that every object you create keeps the expected backend identity.

In [4]:

active_backend = wt.get_backend()
sample = wt.TimeSeries([3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0], dt=1.0, backend=active_backend)
coeffs = sample.to_wdm(nt=4, backend=active_backend)

print(f"Active backend: {active_backend.name}")
print(f"Sample backend: {sample.backend.name}")
print(f"WDM backend: {coeffs.backend.name}")

assert sample.backend is active_backend
assert coeffs.backend is active_backend

active_backend = wt.get_backend() sample = wt.TimeSeries([3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0], dt=1.0, backend=active_backend) coeffs = sample.to_wdm(nt=4, backend=active_backend) print(f"Active backend: {active_backend.name}") print(f"Sample backend: {sample.backend.name}") print(f"WDM backend: {coeffs.backend.name}") assert sample.backend is active_backend assert coeffs.backend is active_backend

Active backend: numpy
Sample backend: numpy
WDM backend: numpy

In [5]:

jax_backend = wt.get_backend("jax")
jax_series = wt.TimeSeries([1.0, 2.0, 3.0, 4.0], dt=0.5, backend=jax_backend)
print(f"Requested backend: {jax_backend.name}")
print(f"JAX series backend: {jax_series.backend.name}")
print(f"JAX array type: {type(jax_series.data).__module__}")
assert jax_series.backend is jax_backend

jax_backend = wt.get_backend("jax") jax_series = wt.TimeSeries([1.0, 2.0, 3.0, 4.0], dt=0.5, backend=jax_backend) print(f"Requested backend: {jax_backend.name}") print(f"JAX series backend: {jax_series.backend.name}") print(f"JAX array type: {type(jax_series.data).__module__}") assert jax_series.backend is jax_backend

Requested backend: jax
JAX series backend: jax
JAX array type: jaxlib._jax

Expected output snapshot¶

On a standard NumPy-only installation, the output should look like this:

Default backend: numpy
Series backend: numpy
Underlying array type: numpy
'numpy' is available: numpy
'jax' is not available in this environment
'cupy' is not available in this environment
Environment default: numpy
Object confirms it is on: numpy
Active backend: numpy
Sample backend: numpy
WDM backend: numpy

On a pip install wdm-transform[jax] installation, the explicit JAX check should look like this:

Requested backend: jax
JAX series backend: jax
JAX array type: jax.numpy

The default backend remains numpy unless you explicitly request jax or set WDM_BACKEND=jax.