jax.numpy dtypes deeply violate the principle of least surprise

[davidmarttila]

Hello,

I didn't open this as an issue since it is not necessarily about a bug in JAX, but about a design decision that I believe should not be carried over into an eventual 1.0 release.

Specifically, it is about JAX dtypes pretending to be numpy dtypes by emulating their hash:

jax/jax/_src/numpy/scalar_types.py

Lines 37 to 41 in 11befd4

	class _ScalarMeta(type):
	dtype: np.dtype
	def __hash__(self) -> int:
	return hash(self.dtype.type)

Overriding the hash to act as if they are the same breaks all kinds of expectations and assumptions made by Python itself and third-party typecheckers, leading to all kinds of completely unintuitive behavior. For example, the python docs state:

Union[X, Y] is equivalent to X | Y and means either X or Y.

jax.numpy dtypes violate both parts of this statement ("equivalent to X | Y", and "either X or Y") as a result of their hash-hacking.

Since Python thinks they are the same type, the | operator deduplicates them to whichever type is mentioned first. This breaks the commutativity of the operator:

>>> type[jnp.float32] | type[np.float32]
type[jax.numpy.float32]
>>> type[np.float32] | type[jnp.float32]
type[numpy.float32]

And it breaks the equivalency of typing.Union and the | operator, since Python seems to perform some form of caching based on what hashes are meant to express:

>>> Union[type[jnp.float32], type[np.float32]]
type[jax.numpy.float32]
>>> Union[type[np.float32], type[jnp.float32]]
type[jax.numpy.float32] # different result from type[np.float32] | type[jnp.float32] !!

Here, the behavior even depends on the ordering of the statements; if the two lines are switched, both reduce to numpy.float32 instead.

Now this would maybe be not so bad if jnp.float32 and np.float32 were actually the same, but they are fundamentally not. JAX dtypes aren't numpy dtypes, not even the standard ones like float32. They don't have the same MRO, they are instances of different metaclasses, and they respond differently to different checks.

So the hash-hacking completely breaks runtime typechecking. For example, with typeguard:

>>> from typeguard import check_type
>>> check_type(np.float32, type[np.float32] | type[jnp.float32])
<class 'numpy.float32'> # check passed
>>> check_type(np.float32, type[jnp.float32] | type[np.float32])
Traceback (most recent call last):
[...]
typeguard.TypeCheckError: class numpy.float32 is not a subclass of jax.numpy.float32

With beartype:

>>> from beartype.door import die_if_unbearable
>>> die_if_unbearable(jnp.float32, type[jnp.float32]) # passes
>>> die_if_unbearable(np.float32, type[np.float32]) # fails because jnp.float32 has been cached based on the hash!
Traceback (most recent call last):
[...] 
value <class 'numpy.float32'> violates type hint type[numpy.float32],
but violation factory get_hint_object_violation() erroneously suggests this object satisfies this hint
[...]

And even other basic checks:

>>> issubclass(np.float32, np.floating)
True
>>> issubclass(jnp.float32, np.floating)
False
>>> jnp.float32 == np.float32
True # this kind of directly contradicts the above, no?

I'm aware that jnp.issubdtype exists, but this behavior just doesn't make sense.

More context and examples on the downstream consequences of this choice is contained in this discussion thread.

[jakevdp]

Unfortunately, this was a design decision made for us long before JAX existed.

np.float32 and friends are not dtypes, rather they are scalar types / constructors that duck-type as dtypes. So in order to implement the NumPy API in jax.numpy, we needed jnp.float32 and friends to be scalar constructors that duck-type as dtypes.

It gets a bit more complicated when you realize that unlike NumPy, JAX doesn't have special types for scalars, but rather represents scalars as zero-dimensional arrays. So jnp.float32 when called needs to return a zero-dimensional jax.Array with dtype float32, and jnp.float32 must also duck-type as a dtype, and also be recognized by APIs like jnp.issubdtype, which requires jnp.float32 and friends to be actual classes (not just functions) with particular equality semantics.

The result of all these requirements is the surprising implementation details you bring up – we could remove those surprises, but it would break JAX's equivalence with NumPy APIs. That would be painful and confusing enough for users that I can't see us ever going down that route.

[davidmarttila]

Thanks for the reply.

I understand that np.float32 and jnp.float32 are types themselves. I understand why JAX needs to replicate this behavior. I don't understand why this requires JAX to override the hashes in a way that breaks the Python specification itself? Because the numpy API by itself does not introduce this problem, it is only introduced by JAX forcing hash(jnp.float32) == hash(np.float32).

This would be less of an issue if JAX actually simply copied over the way that numpy designs dtypes, but there are quite a few differences. np.float* have an MRO hierarchy:

>>> np.float32.__mro__
(<class 'numpy.float32'>, <class 'numpy.floating'>, <class 'numpy.inexact'>, <class 'numpy.number'>, <class 'numpy.generic'>, <class 'object'>)

jnp.float* does not:

>>> jnp.float32.__mro__
(<class 'jax.numpy.float32'>, <class 'object'>)

And extended dtypes that JAX introduces that don't exist in numpy don't fit into the type hierarchy either, even though their types are explicitly instances of np._DTypeMeta:

>>> f16, bf16 = jnp.dtype(jnp.float16), jnp.dtype(jnp.bfloat16)
>>> type(f16)
<class 'numpy.dtypes.Float16DType'>
>>> type(type(f16))
<class 'numpy._DTypeMeta'>
>>> type(f16).__mro__
(<class 'numpy.dtypes.Float16DType'>, <class 'numpy.dtypes._FloatAbstractDType'>, <class 'numpy.dtype'>, <class 'object'>)
>>> type(bf16)
<class 'numpy.dtype[bfloat16]'>
>>> type(type(bf16))
<class 'numpy._DTypeMeta'>
>>> type(bf16).__mro__
 # why does this inherit `numpy.dtype` but not `_FloatAbstractDType`??
(<class 'numpy.dtype[bfloat16]'>, <class 'numpy.dtype'>, <class 'object'>)

[jakevdp]

This would be less of an issue if JAX actually simply copied over the way that numpy designs dtypes, but there are quite a few differences. np.float* have an MRO hierarchy:

I think it's possible this could work – I wasn't around when this code was first written, and I'm not sure whether this was considered. I haven't tried defining the classes this way, but it might be worth trying. Would you like to put that together in a PR and we can see which, if any, tests are affected?

And extended dtypes that JAX introduces that don't exist in numpy don't fit into the type hierarchy either, even though their types are explicitly instances of np._DTypeMeta:

Unfortunately, NumPy's design choices forced our hand on this. We initially tried making bfloat16 and others full participants in the NumPy scalar type hierarchy, but there are many places where NumPy hard-codes the assumption that floating point types are uniquely determined by their bit width. Introducing a 16-bit float that was not float16 led to collisions and breakages, not just for bfloat16, but for float16 itself! Numpy has recently redesigned its dtype system and this may have alleviated the issues, but porting all of ml_dtypes from the old-style to new-style dtypes is a big task, so we haven't had the opportunity to try it yet.

[davidmarttila]

but there are many places where NumPy hard-codes the assumption that floating point types are uniquely determined by their bit width. Introducing a 16-bit float that was not float16 led to collisions and breakages, not just for bfloat16, but for float16 itself!

I see, thanks for the explanation - that is indeed annoying, but I think it doesn't stand in the way of fixing the issues in my top-level post, which are a result of jnp.{bool, uint*, float*, complex*} pretending to be their existing numpy equivalent by emulating their hash, but without replicating their full behavior as regards MRO and subclass checks.

Would you like to put that together in a PR and we can see which, if any, tests are affected?

I just had a look into this and it seems that it's a pretty small modification that doesn't break existing tests and fixes the issues pointed out above -- I'll submit the PR now.

[hawkinsp]

You've partially misunderstood the problem. Hashes, by definition do not have to be unique. So what if np.float32 and jnp.float32 hash equally? That's completely fine, that's what hashes do. The rule is that if two things compare as equal, they should also hash as equal. But the converse is not true. Two unequal things may have the same hash.

I think your issue is with __eq__, not __hash__. We define __eq__ on the JAX ScalarMeta objects to be equal if the underlying dtype would compare as equal, and that's what's surprising you.

i.e., currently we have:

In [1]: import jax.numpy as jnp, numpy as np, itertools

In [2]: for x, y in itertools.permutations([np.int32, np.dtype(np.int32), jnp.int32], 2): print(repr(x), repr(y), x == y)
<class 'numpy.int32'> dtype('int32') True
<class 'numpy.int32'> <class 'jax.numpy.int32'> True
dtype('int32') <class 'numpy.int32'> True
dtype('int32') <class 'jax.numpy.int32'> True
<class 'jax.numpy.int32'> <class 'numpy.int32'> True
<class 'jax.numpy.int32'> dtype('int32') True

(jnp.dtype(jnp.int32) and np.dtype(np.int32) are the same object which is why I didn't include jnp.dtype(jnp.int32)).

Note that np.dtype and np.generic instances already have confusing equality semantics, because they compare as equal but don't have the same hash value. Note that this is a violation of the Python data model: "The only required property is that objects which compare equal have the same hash value" (https://docs.python.org/3/reference/datamodel.html#object.__hash__).

Now, we could solve your problem by removing the equality between np.int32 and jnp.int32.

It's not a hard change to JAX to get this instead:

In [19]: for x, y in itertools.permutations([np.int32, np.dtype(np.int32), jnp.int32], 2): print(repr(x), repr(y), x == y)
<class 'numpy.int32'> dtype('int32') True
<class 'numpy.int32'> <class 'jax.numpy.int32'> False
dtype('int32') <class 'numpy.int32'> True
dtype('int32') <class 'jax.numpy.int32'> True
<class 'jax.numpy.int32'> <class 'numpy.int32'> False
<class 'jax.numpy.int32'> dtype('int32') True

and then your problem would be solved. All one has to do to make this happen is to remove the __eq__, __hash__, and __instancecheck__ methods. The dtype still compares as equal with jnp.int32 because np.dtype looks for a .dtype attribute on its counterpart. And that's highly confusing but also fundamental to how NumPy dtypes work, so that's not changing any time soon.

The main open question is: how much user code does this break? That's really hard to answer without trying it, but my guess is "probably lots". But hey, I can try.

[davidmarttila]

Sure, I agree that the issues would go away instantly if jnp.float32 != np.float32, and then the hashes wouldn't matter. I guess the fundamental issue is the current implementation violating clsA == clsB $\rightarrow$ issubclass(clsA, clsB) and issubclass(clsB, clsA).

I assumed that jnp.float32 == np.float32 was desired behavior because just like you said, there's probably quite some code out there that relies on that. So my PR was aimed at making the consequent true, so to say. But if you are open to having the antecedent be false instead, then this would absolutely solve these issues and be an overall neater fix.

Note that np.dtype and np.generic instances already have confusing equality semantics, because they compare as equal but don't have the same hash value. Note that this is a violation of the Python data model

If JAX would behave the same, ie. jnp.float32 == np.float32 but hash(jnp.float32) != np.float32, then the issues in my original post wouldn't actually exist, since they are caused by Python itself and external typecheckers caching types in dicts and sets. I'm not saying that JAX should behave this way, but I do think it's worth pointing out that it makes sense to talk about hashes when the issue is caching.