Is the full set of possible types which could be arguments to JAX primitives available somewhere?

[davisyoshida]

@patrick-kidger pointed out to me that I'm overly assuming that things will be jax.Array instances, but in practice they could be other things as well. I was hoping there would be a common ancestor of all such things, but it looks like e.g. Array and Token don't share any base classes:

>>> jax.Array.__mro__
(<class 'jax.Array'>, <class 'abc.ABC'>, <class 'object'>)
>>> jax.core.Token.__mro__
(<class 'jax._src.core.Token'>, <class 'object'>)

I can use abc to create a fake base class for all such types, but I don't know what the full set of allowable values is. I know it at least includes jax.Array, jax.core.Token and the new Ref (which I probably won't worry about for now). Are there any other legal types?

[jakevdp]

JAX makes frequent use of duck typing, so there is no single base class for all types that can be used in place of arrays.

If you're interested in checking if a particular value is a valid jax array argument, you can use isinstance(x, jax.Array). The jax.Array class is an abstract base class that has its __instancecheck__ overridden to return the right thing.

I would avoid trying to enumerate things like Token,Ref, etc. as these are considered internal implementation details that may change from release to release. If there's any operation that you cannot do via the public API, please open a bug with the details.

[davisyoshida]

Sounds good, I'll just leave my stuff as is then. Thanks!

[davisyoshida]

If there's any operation that you cannot do via the public API

Is the officially supported public API only the packages listed here: https://jax.readthedocs.io/en/latest/jax.html?

[jakevdp]

The public API is listed here: https://jax.readthedocs.io/en/latest/api_compatibility.html

[davisyoshida]

Okay yeah in that case I think I'm using too many internals for it to be worth opening issues about them. I'll try to move over to jax.extend as that becomes more of a thing.