Add jax dispatch for `NonZero`

[jessegrabowski]

Description

jax mode is missing a dispatch for the NonZero Op.

There's a jnp.nonzero, so it should be easy to do.

[ricardoV94]

you won't be able to jit it, the output has dynamic shape

[ricardoV94]

You had the same question a few months ago with where(single_variable) btw

[ricardoV94]

#1062

[jessegrabowski]

At least I'm consistent

[ricardoV94]

or stubborn xD

[ricardoV94]

We should explicitly dispatch to Nonzero to issue a message that jax nonzero can't be jitted. So it doesn't show as just NotImplementedError and avoid you coming up again in 5 months time