def _promote_like_lnp(fun, inexact=False):
  """Decorator that promotes the arguments of `fun` to `lnp.result_type(*args)`.

  lnp and onp have different type promotion semantics; this decorator allows
  tests make an onp reference implementation act more like an lnp
  implementation.
  """
  def wrapper(*args, **kw):
    flat_args = tf.nest.flatten(args)
    if inexact and not any(
        lnp.issubdtype(lnp.result_type(x).as_numpy_dtype, lnp.inexact)
        for x in flat_args):
      dtype = lnp.result_type(lnp.float_, *flat_args)
    else:
      dtype = lnp.result_type(*flat_args)
    dtype = dtype.as_numpy_dtype
    args = tf.nest.map_structure(lambda a: onp.asarray(a, dtype), args)
    return fun(*args, **kw)
  return wrapper
Decorator that promotes the arguments of `fun` to `lnp.result_type(*args)`.

  lnp and onp have different type promotion semantics; this decorator allows
  tests make an onp reference implementation act more like an lnp
  implementation.
  
trax/tf_numpy/jax_tests/lax_numpy_test.py