jnp.zeros does not support float0 as a dtype

[danieldjohnson]

After #4039, it's no longer possible to pass an integer zero tangent to jax.jvp, it instead needs to be a float0 tangent. However, jnp.zeros doesn't actually support creating arrays of dtype float0, giving:

TypeError: JAX only supports number and bool dtypes, got dtype [('float0', 'V')]

Using np.zeros works fine, but it seems like jnp.zeros should work here too, especially since float0 is a jax-specific thing.

(For context, I have some code that looks like

def some_function_custom_jvp_rule(primals, tangents):
    x, y = primals
    dx, _ = tangents
    _, f_jvp_x = jax.jvp(lambda x: foo(x, y), (x,), (dx,))
    # ... do something with f_jvp_x

and sometimes x is a pytree containing integers. Before, this just worked; now, I have to go in and change the dtypes, because custom_jvp gives me integer tangents but jax.jvp expects float0 tangents. Perhaps this is a bug too? I'm not sure.)

[mattjj]

cc @LenaMartens just as FYI (no obligation implied!)

[NeilGirdhar]

This appears to still be broken despite #4500?

from jax import jvp

def f(x, y):
    return x

print(jvp(f, (3.0, 5), (1.0, 0)))  #  tangent dtype must be float0.Got primal dtype int32; works with float0

[jakevdp]

I think #4500 was intended to fix a different issue. For integer tangents, JVP expects a float0, not a literal zero:

import numpy as np
from jax import dtypes
from jax import jvp

def f(x, y):
    return x

zero = np.zeros((), dtype=dtypes.float0)
print(jvp(f, (3.0, 5), (1.0, zero)))

I don't think there is any plan to change this.

The bug tracked by this issue still stands, i..e. the following fails:

jnp.zeros((), dtype=dtypes.float0)
# TypeError: JAX only supports number and bool dtypes, got dtype [('float0', 'V')] in zeros

[NeilGirdhar]

@jakevdp Okay, that makes sense, thanks for clarifying!

[jakevdp]

I think this is actually working as intended. Even trying to use jax.device_put on a float0 array explicitly passes the ndarray through: https://github.com/google/jax/blob/e32373c3eaa2d8d49372cdfb8698ee2d7af162ab/jax/_src/dispatch.py#L985-L989

I think this was chosen because float0 is not a valid XLA type, perhaps? In any case It seems like we can close this issue.

[wjakob]

I just ran into this -- it seems somehow unnatural to construct NumPy arrays in a CUDA computation so that one use core functionality (jax.vjp, jax.jvp). It would be great if a true fix to jnp.zeros could be considered.

[jakevdp]

Hi @wjakob - can you say more about the issue you're facing? float0 arrays are purely a trace-time entity, and can only exist on the host (XLA does not have any equivalent of a float0 type). Can you say more about what you mean by "a true fix to jnp.zeros"?