"""Backend utilities for ``zea``.
.. note::
Most tensor operations are handled by Keras 3. This module only wraps the
features that Keras does not expose directly: JIT compilation, automatic
differentiation, and device placement.
Public API
----------
:func:`jit`
Unified JIT compilation for JAX (``jax.jit``), TensorFlow
(``tf.function``) and PyTorch (``torch.compile``).
:class:`device`
Context manager that pins all Keras ops to a specific device.
Re-exported as :func:`zea.device`.
:func:`func_on_device`
Run a callable with its tensor arguments moved to a target device.
For ``torch`` this also calls ``.to(device)`` on every input tensor.
:func:`str_to_jax_device`
Resolve a zea device string to the ``jax.Device`` it refers to, for the
rare cases where a JAX device object is needed instead of a string.
:class:`AutoGrad`
Backend-agnostic automatic differentiation wrapper.
"""
import functools
from contextlib import nullcontext
import keras
from zea import log
def _import_tf(force=False):
if not force and keras.backend.backend() != "tensorflow":
return None
try:
import tensorflow as tf
return tf
except ImportError:
return None
def _import_jax(force=False):
if not force and keras.backend.backend() != "jax":
return None
try:
import jax
return jax
except ImportError:
return None
def _import_torch(force=False):
if not force and keras.backend.backend() != "torch":
return None
try:
import torch
return torch
except ImportError:
return None
def _get_backend():
try:
backend_result = keras.backend.backend()
if isinstance(backend_result, str):
return backend_result
else:
# to handle mocked backends during testing
return None
except Exception:
return None
backend = _get_backend()
tf_mod = _import_tf()
jax_mod = _import_jax()
torch_mod = _import_torch()
[docs]
def tf_function(func=None, jit_compile=False, **kwargs):
"""Applies default tf.function to the given function. Only in TensorFlow backend."""
return jit(func, jax=False, torch=False, jit_compile=jit_compile, **kwargs)
[docs]
def jit(func=None, jax=True, tensorflow=True, torch=True, **kwargs):
"""
Applies JIT compilation to the given function based on the current Keras backend.
Can be used as a decorator or as a function.
Args:
func (callable): The function to be JIT compiled.
jax (bool): Whether to enable JIT compilation in the JAX backend.
tensorflow (bool): Whether to enable JIT compilation in the TensorFlow backend.
torch (bool): Whether to enable JIT compilation in the PyTorch backend.
**kwargs: Keyword arguments to be passed to the JIT compiler. One keyword is
handled by zea itself: ``default_compiler_options`` (dict), XLA options
that JAX applies to this function only, and only those the installed
jaxlib supports; an explicit ``compiler_options`` entry overrides them.
Other backends ignore it.
Returns:
callable: The JIT-compiled function.
"""
if func is None:
def decorator(func):
return _jit_compile(func, jax=jax, tensorflow=tensorflow, torch=torch, **kwargs)
return decorator
else:
return _jit_compile(func, jax=jax, tensorflow=tensorflow, torch=torch, **kwargs)
_jit_not_supported_warned = False
def _warn_jit_not_supported(backend_name: str) -> None:
"""Emit the 'JIT not supported' warning at most once, and only when it will
actually be shown (i.e. not when suppressed by log.set_level("ERROR"))."""
global _jit_not_supported_warned
import logging
if not _jit_not_supported_warned and log.logger.isEnabledFor(logging.WARNING):
_jit_not_supported_warned = True
log.warning(
f"JIT compilation not currently supported for backend {backend_name}. "
"Supported backends are 'tensorflow', 'jax' and 'torch'. "
"Initialize zea.Pipeline with jit_options=None to suppress this warning. "
"Falling back to non-compiled mode."
)
def _jit_compile(func, jax=True, tensorflow=True, torch=True, **kwargs):
backend = keras.backend.backend()
if backend == "tensorflow" and tensorflow:
if tf_mod is None:
raise ImportError("TensorFlow is not installed. Please install it to use this backend.")
kwargs.pop("default_compiler_options", None)
jit_compile = kwargs.pop("jit_compile", True)
return tf_mod.function(func, jit_compile=jit_compile, **kwargs)
elif backend == "jax" and jax:
if jax_mod is None:
raise ImportError("JAX is not installed. Please install it to use this backend.")
default_compiler_options = kwargs.pop("default_compiler_options", None)
if default_compiler_options or kwargs.get("compiler_options"):
return _jax_jit_with_compiler_options(func, default_compiler_options, kwargs)
return jax_mod.jit(func, **kwargs)
elif backend == "torch" and torch:
if torch_mod is None:
raise ImportError("PyTorch is not installed. Please install it to use this backend.")
kwargs.pop("default_compiler_options", None)
return torch_mod.compile(func, **kwargs)
elif backend in ("tensorflow", "jax", "torch"):
return func
else:
# Return a lazy wrapper that warns only on first invocation.
# Deferring to call-time lets outer pipelines propagate jit_options=None
# and replace this wrapper before it is ever executed, so no warning
# fires when the user correctly suppresses JIT at the outer level.
_backend = backend
@functools.wraps(func)
def _warn_on_first_call(*args, **kw):
_warn_jit_not_supported(_backend)
return func(*args, **kw)
return _warn_on_first_call
@functools.lru_cache(maxsize=None)
def _jax_compiler_option_supported(name, value) -> bool:
"""Whether the installed jaxlib accepts XLA option ``name`` as a per-function option.
XLA rejects an unknown option at compile time, and experimental flags come and go
between jaxlib releases, so this compiles a trivial function to find out.
"""
try:
jax_mod.jit(lambda: 0, compiler_options={name: value}).lower().compile()
except Exception: # TypeError on a jax.jit without compiler_options, else XLA's error
log.debug(f"XLA compiler option {name!r} is not supported by this jaxlib; skipping it.")
return False
return True
def _jax_jit_with_compiler_options(func, default_compiler_options, jit_kwargs):
"""``jax.jit`` with XLA ``compiler_options``, falling back to none where JAX refuses.
``default_compiler_options`` are only used if jaxlib supports them, and explicit
``compiler_options`` in ``jit_kwargs`` override them.
JAX only accepts ``compiler_options`` on a top-level ``jax.jit``. Called inside
another trace, for example a pipeline inside the user's own ``jax.jit``, the function
is compiled as part of the outer function, whose options then apply; a warning
explains how to pass them there instead.
The support check compiles a function, which would start the JAX backend. That has
to wait until the first call: a pipeline built before :func:`zea.init_device` would
otherwise fix the visible devices before they are selected.
"""
jitted = None # (top_level, nested, compiler_options), built on the first call
warned = False
def build():
compiler_options = {
name: value
for name, value in (default_compiler_options or {}).items()
if _jax_compiler_option_supported(name, value)
}
compiler_options.update(jit_kwargs.get("compiler_options") or {})
# Leave the keyword out when empty: jax.jit only has it from JAX 0.4.36.
options = {k: v for k, v in jit_kwargs.items() if k != "compiler_options"}
nested = jax_mod.jit(func, **options)
if not compiler_options:
return nested, nested, compiler_options
top_level = jax_mod.jit(func, **options, compiler_options=compiler_options)
return top_level, nested, compiler_options
@functools.wraps(func)
def call(*args, **kwargs):
nonlocal jitted, warned
if jitted is None:
jitted = build()
top_level, nested, compiler_options = jitted
if top_level is nested:
return nested(*args, **kwargs)
try:
return top_level(*args, **kwargs)
except ValueError as exc:
if "compiler_options" not in str(exc):
raise
if not warned:
warned = True
_warn_compiler_options_dropped(func, compiler_options)
return nested(*args, **kwargs)
return call
def _warn_compiler_options_dropped(func, options):
"""Warn that ``options`` are lost because ``func`` is jit as an inner function."""
name = getattr(func, "__qualname__", repr(func))
xla_flags = " ".join(
f"--{k}={str(v).lower() if isinstance(v, bool) else v}" for k, v in options.items()
)
log.warning(
f"{name} is traced inside another jax.jit, which does not allow its XLA compiler_options"
f"{options}; without them compilation can be much slower. Pass them to the outer"
f"function `jax.jit(f, compiler_options={options})`, or set XLA_FLAGS='{xla_flags}' "
f"before importing JAX."
)
[docs]
class device:
"""Context manager to run operations on a specific device, regardless of backend.
Normalises device strings across JAX, TensorFlow, and PyTorch so that
``'gpu:0'``, ``'cuda:0'`` and ``'cpu'`` all work with every backend, then
delegates to :func:`keras.device` which handles the per-backend dispatch.
For the ``torch`` backend, :func:`keras.device` sets Keras's internal
device-tracking state so that tensors created by Keras ops land on the
correct device. Existing input tensors are **not** moved automatically —
use ``pipeline(device=..., **inputs)`` or
:func:`zea.backend.func_on_device` when you also need to relocate
pre-existing tensors.
Args:
device (str): Device string, e.g. ``'cuda:0'``, ``'gpu:0'``, or
``'cpu'``.
Example:
.. code-block:: python
# All backends: tensors created by Keras ops are placed on gpu:0
with zea.device("gpu:0"):
output = pipeline(data=data)
# Per-call device with automatic input-tensor movement (all backends)
output = pipeline(device="gpu:0", data=data)
"""
def __init__(self, device: str | None):
if device is None:
self._context = nullcontext()
else:
normalized = self._normalize(device)
self._context = keras.device(normalized)
@staticmethod
def _normalize(device: str) -> str:
"""Normalize device string before passing to ``keras.device``.
Converts ``cuda:N`` → ``gpu:N`` so the string is backend-agnostic;
``keras.device`` itself then converts ``gpu:N`` → ``cuda:N`` when
running under the ``torch`` backend.
"""
device = device.lower()
if device.startswith("auto:"):
raise ValueError(
f"``zea.device`` does not accept 'auto:N' device strings (got {device!r}). "
"Use zea.init_device('auto:N') first to resolve a concrete device, "
"then pass the returned string (e.g. 'gpu:0') to ``zea.device``."
)
# Normalise to gpu:N; keras.device handles gpu → cuda for the torch backend.
return device.replace("cuda", "gpu")
def __enter__(self):
self._context.__enter__()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self._context.__exit__(exc_type, exc_val, exc_tb)
# Private alias so func_on_device can reference the class without clashing
# with its own `device` parameter.
_DeviceContext = device
[docs]
def str_to_jax_device(device: str):
"""Resolve a device string to the :class:`jax.Device` object it refers to.
Accepts the same strings as :class:`zea.device` (``'gpu:0'``, ``'cuda:0'``,
``'cpu'``, ...) and resolves them against the devices JAX actually sees.
Only needed where a JAX device object is required rather than a string;
prefer :class:`zea.device` for placing ops on a device.
Args:
device (str): Device string, e.g. ``'gpu:0'`` or ``'cpu'``. Without an
index the first device of that type is returned.
Returns:
jax.Device: The corresponding JAX device.
Raises:
ImportError: If JAX is not installed.
ValueError: If the string is malformed, or names a device JAX cannot see.
"""
jax = _import_jax(force=True)
if jax is None:
raise ImportError("JAX is not installed. Please install it to resolve JAX devices.")
if not isinstance(device, str):
raise ValueError(f"Device must be a string, got {type(device)}")
# Share one definition of what a device string means with ``zea.device``.
device_type, _, device_number = _DeviceContext._normalize(device).partition(":")
try:
device_number = int(device_number) if device_number else 0
except ValueError:
raise ValueError(f"Device number in '{device}' is not an integer.") from None
try:
available = jax.devices(device_type)
except RuntimeError as exc:
# jax raises for a device type it cannot initialize (e.g. 'tpu' without
# libtpu). From here that is just an unavailable device, like any other.
raise ValueError(f"No JAX devices available for type '{device_type}': {exc}") from exc
if not available:
raise ValueError(f"No JAX devices available for type '{device_type}'.")
if device_number < 0 or device_number >= len(available):
raise ValueError(f"Device '{device}' is not available; JAX devices found: {available}")
return available[device_number]
[docs]
def func_on_device(func, device, *args, **kwargs):
"""Run ``func`` with all tensor arguments placed on ``device``.
For the ``torch`` backend, every tensor argument is explicitly moved with
``.to(device)`` before the call. For JAX and TensorFlow the function is
executed inside an :class:`zea.backend.device` context, which routes newly created
tensors to the requested device.
Args:
func (callable): Function to call.
device (str): Target device, e.g. ``'cpu'``, ``'gpu:0'``, ``'cuda:1'``.
*args: Positional arguments forwarded to ``func``.
**kwargs: Keyword arguments forwarded to ``func``.
Returns:
Output of ``func(*args, **kwargs)``.
"""
if device is None:
return func(*args, **kwargs)
if keras.backend.backend() == "torch":
import torch
_device = torch.device(device.lower().replace("gpu", "cuda"))
def _move(x):
if isinstance(x, torch.Tensor):
return x.to(_device)
if isinstance(x, (list, tuple)):
return type(x)(_move(i) for i in x)
if isinstance(x, dict):
return {k: _move(v) for k, v in x.items()}
return x
args = _move(args)
kwargs = _move(kwargs)
with _DeviceContext(device):
return func(*args, **kwargs)