Actually this problem is more complicated than I though. Could it be that numpy dtypes and jax dtype's are different?

cc @tomhennigan

It's a bit tricky to get these right, since there are several different types referred to as 'dtype'. Without looking at it closely, one guess is we could try onp.dtype(self.dtype).name.

I think the test in api_test.py covers this with jax dtype so maybe try @mattjj suggestion and run those tests?

Thanks @mattjj and @tomhennigan the suggested solution seems to work correctly. I can't really think of another edge case at this point and the tests pass.

I actually think that there is real confusion between numpy.dtype and jax.numpy.dtype. These are not the same. However, ShapedDtypeStruct is created in the code with dtype: numpy.dtype. Hence, the correct way to print the type should be self.dtype.name, as Jonathan first suggested.

There is a bug in api_test.py::test_shape_dtype_struct and test_shape_dtype_struct_scalar whereby ShapedDtypeStruct is initialized with a np.float32. It should be onp.float32. (All the rest of api.test uses onp.float32 !)

BTW, the alternative onp.dtype(self.dtype).name also works, but is just working around the fact that the test initializes the struct with the wrong types.

@jheek When you get a chance, can you please take a look at the latest comments?

Most of the jax api accepts both numpy dtypes and jax dtypes, if I'm correct. Because ShapeDtypeStruct can be constructed by users I think it's best to have a fix that allows both flavours of dtypes to be used.

If we do want to allow flexibility for the constructor, how do you feel about adding the onp.dtype conversion in the constructor; then the usage can depend on a precise type.

@jheek PTAL at my suggestion

In the meantime, I'm using this patch. This may be inconsistent between np.dtype/jnp.dtype but at least it allow to print(jax.eval_shape(fn)) without crashing.

def new_repr(self):
  return f'{type(self).__qualname__}(shape={self.shape}, dtype={self.dtype})'

jax.ShapeDtypeStruct.__repr__ = new_repr
jax.ShapeDtypeStruct.__str__ = new_repr

@gnecula yes I agree that makes more sense and it seems to work correctly.

@gnecula yes I agree that makes more sense and it seems to work correctly.

Have you implemented this, it does not seem to be in the last commit you pushed 7c86cf1

seems like I made a mistake when pushing changes. Should be good now

Fixes #2320