def _promote_like_jnp(fun, inexact=False):
    'Decorator that promotes the arguments of `fun` to `jnp.result_type(*args)`.\n\n  jnp and onp have different type promotion semantics; this decorator allows\n  tests make an onp reference implementation act more like an jnp\n  implementation.\n  '

    def wrapper(*args, **kw):
        flat_args = tree_util.tree_leaves(args)
        if (inexact and (not any((jnp.issubdtype(jnp.result_type(x), jnp.inexact) for x in flat_args)))):
            dtype = jnp.result_type(jnp.float_, *flat_args)
        else:
            dtype = jnp.result_type(*flat_args)
        args = tree_util.tree_map((lambda a: onp.asarray(a, dtype)), args)
        return fun(*args, **kw)
    return wrapper
Decorator that promotes the arguments of `fun` to `jnp.result_type(*args)`.

jnp and onp have different type promotion semantics; this decorator allows
tests make an onp reference implementation act more like an jnp
implementation.
def _promote_like_jnp(fun, inexact=False):
    'Decorator that promotes the arguments of `fun` to `jnp.result_type(*args)`.\n\n  jnp and onp have different type promotion semantics; this decorator allows\n  tests make an onp reference implementation act more like an jnp\n  implementation.\n  '

    def wrapper(*args, **kw):
        flat_args = tree_util.tree_leaves(args)
        if (inexact and (not any((jnp.issubdtype(jnp.result_type(x), jnp.inexact) for x in flat_args)))):
            dtype = jnp.result_type(jnp.float_, *flat_args)
        else:
            dtype = jnp.result_type(*flat_args)
        args = tree_util.tree_map((lambda a: onp.asarray(a, dtype)), args)
        return fun(*args, **kw)
    return wrapper