EVOLUTION-MANAGER

Edit File: jax.py

"""
Required functions for optimized contractions of numpy arrays using jax.
"""

import numpy as np

from ..sharing import to_backend_cache_wrap

__all__ = ["build_expression", "evaluate_constants"]

_JAX = None

def _get_jax_and_to_jax():
    global _JAX
    if _JAX is None:
        import jax

@to_backend_cache_wrap
        @jax.jit
        def to_jax(x):
            return x

_JAX = jax, to_jax

return _JAX

def build_expression(_, expr):  # pragma: no cover
    """Build a jax function based on ``arrays`` and ``expr``.
    """
    jax, _ = _get_jax_and_to_jax()

jax_expr = jax.jit(expr._contract)

def jax_contract(*arrays):
        return np.asarray(jax_expr(arrays))

return jax_contract

def evaluate_constants(const_arrays, expr):  # pragma: no cover
    """Convert constant arguments to jax arrays, and perform any possible
    constant contractions.
    """
    jax, to_jax = _get_jax_and_to_jax()

return expr(*[to_jax(x) for x in const_arrays], backend='jax', evaluate_constants=True)