Zero class

[mariogeiger]

I'm implementing a class Zero that would represent an array of zeros.
The aim is to simplify operations like that:

Zero * ndarray -> Zero          # no operation, just a broadcast
Zero + ndarray -> ndarray       # no operation, just a broadcast
einsum(..., Zero, ...) -> Zero  # no operation, just a broadcast

I guess a lot of people have already thought about this, is there already a solution?
(I saw that ad_util.Zero exists but I didn't understand if it is for this purpose.)

[jakevdp]

Hi - I don't think there's any easy way to do this – JAX doesn't provide any hooks within jax.numpy operations to override behavior on custom array-like objects. Regarding ad_util.Zero: it serves a different purposel; it is a placeholder for gradients of integers within automatic differentiation, and is not really handled outside the context of autodiff.

One fairly heavy-weight way you might approach a problem like this is by implementing a custom JAX transformation that would convert a JAX function into an equivalent function that accepts your custom Zero class. This would require defining Zero-aware implementations of each of JAX's primitive operations.

[mariogeiger]

Oh gosh I see the problems now... Zero is not handled by anything otherwise (vmap, grad, jit...).
First I thought I would be fine if Zero is only used internally but then I cannot have a function that inputs/outputs a Zero and is vmaped.

[PhilipVinc]

This makes me wish Jax.numpy was built with multiple dispatch…

[froystig]

You can register the class as a pytree in order to define its behavior as input/output to transformations like vmap, but @jakevdp's point about overloading numpy remains.