`jnp.isposinf` and `float4_e2m1fn` don't play well together

[gonnet]

According to the docs (I'm using v0.5.4), float4_e2m1fn does not support Inf values, and the following

jnp.astype(jnp.inf, jnp.float4_e2m1fn)

produces, as expected, the rounded-down value Array(6, dtype=float4_e2m1fn).

However, if I feed this value to jnp.isposinf, e.g. as a test whether a number format supports +Inf values, i.e.

jnp.isposinf(jnp.astype(jnp.inf, jnp.float4_e2m1fn))

I get the following output:

Array(True, dtype=bool)

which is wrong.

As a workaround, is there any other way to programmatically figure out if a jnp.dtype can represent non-finite values such as Inf or NaN?

Cheers, Pedro

[jakevdp]

Thanks for the report! we can definitely fix this. In the meantime, you could use a function like this one to check whether an input supports inf:

In [1]: import numpy as np

In [2]: import ml_dtypes

In [3]: def can_represent_inf(dtype):
   ...:     return np.isinf(np.array(np.inf, dtype).astype('float64'))
   ...: 

In [4]: can_represent_inf(ml_dtypes.float8_e3m4)
Out[4]: np.True_

In [5]: can_represent_inf(ml_dtypes.float4_e2m1fn)
Out[5]: np.False_

[gonnet]

Cool, thanks @jakevdp for the workaround and quick fix!