def onp_fun(a, b):
    a = (a if (lhs_dtype != jnp.bfloat16) else a.astype(onp.float32))
    b = (b if (rhs_dtype != jnp.bfloat16) else b.astype(onp.float32))
    dtype = jnp.promote_types(lhs_dtype, rhs_dtype)
    return onp.tensordot(a, b, axes).astype(dtype)
def onp_fun(a, b):
    a = (a if (lhs_dtype != jnp.bfloat16) else a.astype(onp.float32))
    b = (b if (rhs_dtype != jnp.bfloat16) else b.astype(onp.float32))
    dtype = jnp.promote_types(lhs_dtype, rhs_dtype)
    return onp.tensordot(a, b, axes).astype(dtype)