Repository: google/jax

Issue:
Various type inconsistencies
Please consider adding `jax.numpy.generic` as an alias for `numpy.generic` (it seems to be missing).  Also,
```
In [5]: jnp.issubdtype(jnp.bfloat16, jnp.inexact)
Out[5]: True

In [6]: issubclass(jnp.bfloat16, jnp.inexact)
Out[6]: False

In [7]: issubclass(np.float32, np.inexact)
Out[7]: True
```
Line 6 should probably return true, which means somehow registering the types like numpy does.


Discussion:
I'm not sure that would be correct; for example:
```python
import numpy as np
print(type(np.float16(0)))
# numpy.float16

import jax.numpy as jnp
type(jnp.float16(0))
# jaxlib.xla_extension.DeviceArray
```
JAX does not return an instance of a `float16` scalar in this case, but rather returns a zero-dimensional JAX array. THis choice to unify arrays and scalars in JAX is a deliberate one: we have specifically decided to *not* honor the numpy Python type hierarchy in the case of scalar-like instantiations, so enforcing that `jnp.float16` is a subclass of `jnp.inexact` would be misleading.

What do you think?
That's a great point.  How am I supposed to say that `T` is one of the inexact types?  In Numpy, I say: `T: type[np.inexact]`.  The corresponding thing doesn't work in Jax.
I think you're looking for `jnp.issubdtype(t, jnp.inexact)`. This should handle `bfloat16` as well as the standard numpy types.
Yes, I know how to check it programmatically.  I'm asking how I should annotate it.
Sorry, I just noticed that your question is more a question of type annotation than in-program type checking (that wasn't clear in your initial question).

Setting JAX aside, I don't think using `T: type[np.inexact]` is the best approach, becuase (for example) `np.dtype('float64')` will fail this.

Instead, you should probably use the new [`DtypeLike`](https://numpy.org/devdocs/reference/typing.html#dtypelike) typedef in newer numpy versions. Unfortunately, this doesn't give any mechanism to limit inputs to inexact types, but I think that Python type checking simply doesn't have the granularity to do that correctly.

It's one of the reasons that JAX internally aliases so many types to `Any`: Python type declaration is just generally not expressive enough to work properly with array computing code, except at a very superficial level.
> Python type declaration is just generally not expressive enough to work properly with array computing code, except at a very superficial level.

I agree with this.  I'm just trying to do the best I can.

Anyway, for this issue, you've convinced me.  I guess all that's left is adding `jnp.generic`?

> Anyway, for this issue, you've convinced me. I guess all that's left is adding jnp.generic?

Sure – I'll take care of that
https://github.com/google/jax/pull/9471

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -424,6 +424,7 @@ def _make_scalar_type(np_scalar_type):
 float_ = float32 if dtypes.float_ == np.float32 else float64
 complex_ = complex64 if dtypes.complex_ == np.complex64 else complex128
 
+generic = np.generic
 number = np.number
 inexact = np.inexact
 complexfloating = np.complexfloating
diff --git a/jax/numpy/__init__.py b/jax/numpy/__init__.py
--- a/jax/numpy/__init__.py
+++ b/jax/numpy/__init__.py
@@ -160,6 +160,7 @@
     full as full,
     full_like as full_like,
     gcd as gcd,
+    generic as generic,
     geomspace as geomspace,
     get_printoptions as get_printoptions,
     gradient as gradient,


Test patch: