def onp_fun(a, b):
    a = (a.astype(onp.float32) if (lhs_dtype == jnp.bfloat16) else a)
    b = (b.astype(onp.float32) if (rhs_dtype == jnp.bfloat16) else b)
    out = onp.cross(a, b, axisa, axisb, axisc, axis)
    return out.astype(jnp.promote_types(lhs_dtype, rhs_dtype))
def onp_fun(a, b):
    a = (a.astype(onp.float32) if (lhs_dtype == jnp.bfloat16) else a)
    b = (b.astype(onp.float32) if (rhs_dtype == jnp.bfloat16) else b)
    out = onp.cross(a, b, axisa, axisb, axisc, axis)
    return out.astype(jnp.promote_types(lhs_dtype, rhs_dtype))