Semantics of tensor-shape (2, 0, 1)

[pratnali-aws]

I came across this test - https://github.com/google/jax/blob/d0008fbe4ad67de2d7c5a9dccd6a28ab1ef03802/tests/lax_test.py#L1577
and wanted to understand the semantics of a tensor shaped (2, 0, 1). Specifically, what does the '0' dimension refer to and mean?

[jakevdp]

Hi - the semantics of 0 in a dimension mean that the dimension is of size 0. I think you'll find that array operations have consistent semantics in this case. For example:

In [1]: import jax.numpy as jnp

In [2]: x = jnp.zeros((2, 0, 1))

In [3]: x.shape
Out[3]: (2, 0, 1)

In [4]: x.shape[1]
Out[4]: 0

In [5]: x.size  # prod(shape)
Out[5]: 0

In [6]: x.ravel().shape  # ravel() = reshape to (size,)
Out[6]: (0,)

In [7]: x.mT.shape  # mT = matrix transpose
Out[7]: (2, 1, 0)

Does that make sense?

[pratnali-aws]

Thanks for the reply Jake! We want to lower this kind of tensor and will have to teach the compiler to deal with a 0-sized tensors.

Thanks!