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):
        import numpy as np  # type: ignore

        return np.asarray(jax_expr(arrays))

    return jax_contract
Build a jax function based on ``arrays`` and ``expr``.
opt_einsum/backends/jax.py