Broadcasting of size-0 dimensions not implemented

[hawkinsp]

Numpy supports broadcasts with size-0 dimensions against size-1 dimensions:

onp.ones([0,1]) + onp.ones([1,128])

produces:

array([], shape=(0, 128), dtype=float64)

However

to_device = jax.jit(lambda x:x)
to_device(np.ones([0,1])) + to_device(np.ones([1,128]))
ValueError: Incompatible shapes for broadcasting: ((0, 1), (1, 128))

The broadcasting rule computes the output shape as

result_shape = onp.max(shapes, axis=0)

but it probably needs to be something like this:

min_shape = onp.min(shapes, axis=0)
max_shape = onp.max(shapes, axis=0)
result_shape = onp.where(min_shape == 0, 0, max_shape)