FEAT: configure backend precision

benchmarks/precision.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+import sympy as sp
+
+from tensorwaves import configure
+from tensorwaves.function.sympy import create_function
+
+if TYPE_CHECKING:
+    from typing import Literal
+
+    from tensorwaves.config import Precision
+    from tensorwaves.interface import DataSample
+
+    Backend = Literal["jax", "tensorflow"]
+
+
+@pytest.mark.benchmark(group="precision")
+@pytest.mark.parametrize("backend", ["jax", "tensorflow"])
+@pytest.mark.parametrize("precision", ["float32", "float64"])
+def test_precision(benchmark, backend: Backend, precision: Precision) -> None:
+    _configure_backend(backend, precision)
+    x = sp.Symbol("x")
+    function = create_function(sp.sin(x) ** 2 + sp.exp(-x), backend=backend)
+    data = _create_data(backend)
+    result = benchmark(lambda: _evaluate(function, data, backend))
+    _assert_precision(backend, precision, data, result)
+
+
+def _configure_backend(backend: Backend, precision: Precision) -> None:
+    if backend == "jax":
+        configure(jax_precision=precision)
+    else:
+        configure(tensorflow_precision=precision)
+
+
+def _create_data(backend: Backend):
+    if backend == "jax":
+        import jax.numpy as jnp
+
+        return {"x": jnp.linspace(0, 10, num=1_000_000).block_until_ready()}
+
+    import tensorflow.experimental.numpy as tnp  # ty: ignore[unresolved-import]
+
+    return {"x": tnp.linspace(0, 10, num=1_000_000)}
+
+
+def _evaluate(function, data: DataSample, backend: Backend):
+    result = function(data)
+    if backend == "jax":
+        result.block_until_ready()
+    return result
+
+
+def _assert_precision(
+    backend: Backend, precision: Precision, data: dict, result
+) -> None:
+    assert data["x"].dtype.name == precision
+    assert result.dtype.name == precision
+    if backend == "jax":
+        import jax
+
+        assert jax.config.x64_enabled == (precision == "float64")
+    else:
+        import tensorflow.experimental.numpy as tnp  # ty: ignore[unresolved-import]
+
+        assert tnp.asarray(1.0).dtype.name == precision

benchmarks/unbinned_nll.py

@@ -11,6 +11,7 @@
 import tensorflow as tf
 import tensorflow.experimental.numpy as tnp  # ty: ignore[unresolved-import]
 
+from tensorwaves import configure
 from tensorwaves.estimator import UnbinnedNLL
 from tensorwaves.function import ParametrizedBackendFunction
 
@@ -203,8 +204,7 @@ def estimator_samples() -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
 def jax_intensities(
     intensities: tuple[np.ndarray, np.ndarray],
 ) -> tuple[jax.Array, jax.Array]:
-    jax.config.update("jax_enable_x64", True)
-
+    configure(jax_precision="float64")
     data_intensities, phsp_intensities = intensities
     jax_data_intensities = jnp.asarray(data_intensities).block_until_ready()
     jax_phsp_intensities = jnp.asarray(phsp_intensities).block_until_ready()
@@ -223,8 +223,7 @@ def tensorflow_intensities(
 def jax_estimator_samples(
     estimator_samples: tuple[dict[str, np.ndarray], dict[str, np.ndarray]],
 ) -> tuple[dict[str, jax.Array], dict[str, jax.Array]]:
-    jax.config.update("jax_enable_x64", True)
-
+    configure(jax_precision="float64")
     data, phsp = estimator_samples
     return (
         {"x": jnp.asarray(data["x"]).block_until_ready()},

docs/install.md

@@ -60,6 +60,25 @@ pip install tensorwaves[jax,scipy,tf]
 pip install tensorwaves[all]  # all runtime dependencies
 ```
 
+## Backend precision
+
+TensorWaves uses 64-bit precision for JAX and TensorFlow by default. Configure 32-bit precision before creating backend arrays or TensorWaves functions with:
+
+```python
+import tensorwaves
+
+tensorwaves.configure(
+    jax_precision="float32",
+    tensorflow_precision="float32",
+)
+```
+
+The two options are independent and can be specified separately.
+
+Precision is a property of the backend, not of TensorWaves, so it applies from the moment it is set: arrays that were created earlier keep the dtype they were created with. A backend that has already been imported is reconfigured immediately, and one that has not is configured when TensorWaves first uses it. Calling `tensorwaves.configure()` right after your imports therefore covers both cases.
+
+TensorWaves also respects JAX's `JAX_ENABLE_X64` environment variable. An explicit call to `tensorwaves.configure()` takes precedence over it. Note that JAX itself only reads that variable when it is imported, whereas TensorWaves reads it when it first uses JAX. TensorFlow has no equivalent environment variable, so `tensorflow_precision` is the only way to select its precision.
+
 :::::{container} full-width
 
 ::::{dropdown} **GPU support**

src/tensorwaves/__init__.py

@@ -20,10 +20,12 @@
 """
 
 __all__ = [
+    "configure",
     "data",
     "estimator",
     "function",
     "optimizer",
 ]
 
 from . import data, estimator, function, optimizer
+from .config import configure

src/tensorwaves/config.py

@@ -0,0 +1,119 @@
+"""Configure optional computational backends."""
+
+from __future__ import annotations
+
+import os
+import sys
+from dataclasses import dataclass
+from importlib import import_module
+from typing import TYPE_CHECKING, Literal, get_args
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from types import ModuleType
+
+Precision = Literal["float32", "float64"]
+
+
+def configure(
+    *,
+    jax_precision: Precision | None = None,
+    tensorflow_precision: Precision | None = None,
+) -> None:
+    """Set the precision used by computational backends.
+
+    Call this function before creating backend arrays or TensorWaves functions. By
+    default, TensorWaves uses 64-bit precision for JAX and TensorFlow. TensorWaves
+    respects ``JAX_ENABLE_X64`` if ``jax_precision`` is not specified.
+    """
+    _validate_precision("jax_precision", jax_precision)
+    _validate_precision("tensorflow_precision", tensorflow_precision)
+    if jax_precision is not None:
+        _jax_config.precision = jax_precision
+        _configure_imported_module("jax", _set_jax_precision, jax_precision)
+    if tensorflow_precision is not None:
+        _tensorflow_config.precision = tensorflow_precision
+        _configure_imported_module(
+            "tensorflow", _set_tensorflow_precision, tensorflow_precision
+        )
+
+
+def _configure_imported_module(
+    module_name: str,
+    set_precision: Callable[[ModuleType, Precision], None],
+    precision: Precision,
+) -> None:
+    """Set the precision on a backend that has already been imported.
+
+    A backend that has not been imported yet is configured on first use, so that
+    :func:`configure` never triggers a backend import itself.
+    """
+    module = sys.modules.get(module_name)
+    if module is not None:
+        set_precision(module, precision)
+
+
+def _initialize_jax() -> ModuleType:
+    jax = import_module("jax")
+
+    if not _jax_config.initialized:
+        precision = _jax_config.precision
+        if precision is None:
+            precision = _precision_from_flag(os.environ.get("JAX_ENABLE_X64", "1"))
+        _set_jax_precision(jax, precision)
+        _jax_config.initialized = True
+    return jax
+
+
+def _set_jax_precision(jax: ModuleType, precision: Precision) -> None:
+    jax.config.update("jax_enable_x64", precision == "float64")
+
+
+def _precision_from_flag(value: str) -> Precision:
+    """Interpret a JAX-style boolean environment variable value.
+
+    >>> _precision_from_flag("1"), _precision_from_flag("false")
+    ('float64', 'float32')
+    """
+    return "float64" if value.strip().lower() in {"1", "true", "yes"} else "float32"
+
+
+def _initialize_tensorflow() -> ModuleType:
+    tf = import_module("tensorflow")
+
+    if not _tensorflow_config.initialized:
+        _set_tensorflow_precision(tf, _tensorflow_precision())
+        _tensorflow_config.initialized = True
+    return tf
+
+
+def _set_tensorflow_precision(tensorflow: ModuleType, precision: Precision) -> None:
+    tensorflow.experimental.numpy.experimental_enable_numpy_behavior(
+        prefer_float32=precision == "float32"
+    )
+
+
+def _tensorflow_precision() -> Precision:
+    return _tensorflow_config.precision or "float64"
+
+
+def _validate_precision(name: str, precision: object) -> None:
+    if precision is not None and precision not in get_args(Precision):
+        msg = f"{name} must be 'float32', 'float64', or None"
+        raise ValueError(msg)
+
+
+@dataclass
+class _JaxConfig:
+    precision: Precision | None = None
+    initialized: bool = False
+
+
+@dataclass
+class _TensorFlowConfig:
+    precision: Precision | None = None
+    initialized: bool = False
+
+
+_jax_config = _JaxConfig()
+_tensorflow_config = _TensorFlowConfig()

src/tensorwaves/data/phasespace.py

@@ -107,7 +107,7 @@ def __init__(
         except ImportError:  # pragma: no cover
             raise_missing_module_error("phasespace", extras_require="phsp")
         sorted_ids = sorted(final_state_masses)
-        self.__phsp_gen = phasespace.nbody_decay(  # ty:ignore[possibly-unresolved-reference]
+        self.__phsp_gen = phasespace.nbody_decay(
             mass_top=initial_state_mass,
             masses=[final_state_masses[i] for i in sorted_ids],
             names=list(map(str, sorted_ids)),

src/tensorwaves/data/rng.py

@@ -6,6 +6,7 @@
 
 import numpy as np
 
+from tensorwaves.config import _tensorflow_precision
 from tensorwaves.function._backend import raise_missing_module_error
 from tensorwaves.interface import RealNumberGenerator
 
@@ -41,11 +42,11 @@ class TFUniformRealNumberGenerator(RealNumberGenerator):
 
     def __init__(self, seed: int | None = None) -> None:
         try:
-            from tensorflow import float64  # ruff:ignore[import-outside-top-level]
+            import tensorflow as tf  # ruff:ignore[import-outside-top-level]
         except ImportError:  # pragma: no cover
             raise_missing_module_error("tensorflow", extras_require="tf")
         self.seed = seed
-        self.dtype = float64  # ty:ignore[possibly-unresolved-reference]
+        self.dtype = tf.float32 if _tensorflow_precision() == "float32" else tf.float64
 
     def __call__(
         self, size: int, min_value: float = 0.0, max_value: float = 1.0
@@ -78,10 +79,10 @@ def _get_tensorflow_rng(seed: SeedLike | None = None) -> tf.random.Generator:
         raise_missing_module_error("tensorflow", extras_require="tf")
 
     if seed is None:
-        return tf.random.get_global_generator()  # ty:ignore[possibly-unresolved-reference]
+        return tf.random.get_global_generator()
     if isinstance(seed, int):
-        return tf.random.Generator.from_seed(seed=seed)  # ty:ignore[possibly-unresolved-reference]
-    if isinstance(seed, tf.random.Generator):  # ty:ignore[possibly-unresolved-reference]
+        return tf.random.Generator.from_seed(seed=seed)
+    if isinstance(seed, tf.random.Generator):
         return seed
     msg = f"Cannot create a tf.random.Generator from a {type(seed).__name__}"
     raise TypeError(msg)

src/tensorwaves/estimator.py

@@ -7,6 +7,7 @@
 
 from typing import TYPE_CHECKING
 
+from tensorwaves.config import _initialize_jax
 from tensorwaves.data.transform import SympyDataTransformer
 from tensorwaves.function._backend import find_function, raise_missing_module_error
 from tensorwaves.function.sympy import create_parametrized_function, prepare_caching
@@ -80,11 +81,10 @@ def gradient_creator(
 ) -> Callable[[Mapping[str, ParameterValue]], dict[str, ParameterValue]]:
     if backend == "jax":
         try:
-            import jax  # ruff:ignore[import-outside-top-level]
+            jax = _initialize_jax()
         except ImportError:  # pragma: no cover
             raise_missing_module_error("jax", extras_require="jax")
-        jax.config.update("jax_enable_x64", True)  # ty:ignore[possibly-unresolved-reference]
-        gradient = jax.grad(function)  # ty:ignore[possibly-unresolved-reference]
+        gradient = jax.grad(function)
 
         def conjugated_gradient(
             parameters: Mapping[str, ParameterValue],

src/tensorwaves/function/_backend.py

@@ -6,9 +6,11 @@
 from typing import TYPE_CHECKING
 from warnings import warn
 
+from tensorwaves.config import _initialize_jax, _initialize_tensorflow
+
 if TYPE_CHECKING:
     from collections.abc import Callable
-    from typing import ParamSpec, TypeVar
+    from typing import NoReturn, ParamSpec, TypeVar
 
     P = ParamSpec("P")
     T = TypeVar("T")
@@ -40,27 +42,24 @@ def get_backend_modules(backend: str | tuple | dict) -> str | tuple | dict:
     if isinstance(backend, str):
         if backend == "jax":
             try:
-                import jax
+                _initialize_jax()
                 import jax.numpy as jnp
                 import jax.scipy as jsp
             except ImportError:  # pragma: no cover
                 raise_missing_module_error("jax", extras_require="jax")
-            jax.config.update("jax_enable_x64", True)  # ty:ignore[possibly-unresolved-reference]
-            return jnp, jsp.special  # ty:ignore[possibly-unresolved-reference]
+            return jnp, jsp.special
         if backend in {"numpy", "numba"}:
             import numpy as np
 
             return np, np.__dict__
             # returning only np.__dict__ does not work well with conditionals
         if backend in {"tensorflow", "tf"}:
             try:
-                import tensorflow as tf
-                import tensorflow.experimental.numpy as tnp  # ty:ignore[unresolved-import]
-                from tensorflow.python.ops.numpy_ops import np_config
+                tf = _initialize_tensorflow()
+                tnp = tf.experimental.numpy
             except ImportError:  # pragma: no cover
                 raise_missing_module_error("tensorflow", extras_require="tf")
-            np_config.enable_numpy_behavior()  # ty:ignore[possibly-unresolved-reference]
-            return tnp.__dict__, tf  # ty:ignore[possibly-unresolved-reference]
+            return tnp.__dict__, tf
 
     return backend
 
@@ -85,14 +84,14 @@ def jit_compile(backend: str) -> Callable[[Callable[P, T]], Callable[P, T]]:
             import jax
         except ImportError:  # pragma: no cover
             raise_missing_module_error("jax", extras_require="jax")
-        return jax.jit  # ty:ignore[possibly-unresolved-reference]
+        return jax.jit
 
     if backend == "numba":
         try:
             import numba
         except ImportError:  # pragma: no cover
             raise_missing_module_error("numba", extras_require="numba")
-        return partial(numba.jit, forceobj=True, parallel=True)  # ty:ignore[possibly-unresolved-reference]
+        return partial(numba.jit, forceobj=True, parallel=True)
 
     msg = f"Backend {backend} does not yet support JIT compilation"
     warn(msg, category=UserWarning, stacklevel=3)
@@ -103,7 +102,9 @@ def _do_not_compile(function: Callable[P, T]) -> Callable[P, T]:
     return function
 
 
-def raise_missing_module_error(module_name: str, *, extras_require: str = "") -> None:
+def raise_missing_module_error(
+    module_name: str, *, extras_require: str = ""
+) -> NoReturn:
     """Raise an `ImportError` with install instructions.
 
     >>> raise_missing_module_error("missing")