import jax
import jax.numpy as jnp

### ----- modified exerpt from optimistix/_misc.py -----

def _asarray(dtype, x):
    return jnp.asarray(x, dtype=dtype)

# remove custom jvp to resolve issue
_asarray = jax.custom_jvp(_asarray, nondiff_argnums=(0,))

@_asarray.defjvp
def _asarray_jvp(dtype, x, tx):
    (x,) = x
    (tx,) = tx
    return _asarray(dtype, x), _asarray(dtype, tx)

def inexact_asarray(x):
    dtype = jnp.result_type(x)
    if not jnp.issubdtype(jnp.result_type(x), jnp.inexact):
        dtype = jnp.float32
    return _asarray(dtype, x)

### ----- end -----

@jax.jit
def const_eval(y0):
    with jax.ensure_compile_time_eval():
        y_stub = jax.tree.map(lambda a: jnp.zeros_like(a), y0)
        print(y_stub) # 0

        # not greedily evaluated
        out = jax.tree.map(inexact_asarray, y_stub)
        print(out) # JitTracer<float32[]>

        # works as expected when jit'd
        out = jax.jit(lambda a: jax.tree.map(inexact_asarray, a))(y_stub)
        print(out) # 0.0

const_eval(1)