def onp_dot(x, y):
    x = (x.astype(onp.float32) if (lhs_dtype == jnp.bfloat16) else x)
    y = (y.astype(onp.float32) if (rhs_dtype == jnp.bfloat16) else y)
    return onp.dot(x, y).astype(jnp.promote_types(lhs_dtype, rhs_dtype))
def onp_dot(x, y):
    x = (x.astype(onp.float32) if (lhs_dtype == jnp.bfloat16) else x)
    y = (y.astype(onp.float32) if (rhs_dtype == jnp.bfloat16) else y)
    return onp.dot(x, y).astype(jnp.promote_types(lhs_dtype, rhs_dtype))