Bug in `nonzero_idx`?

[krstopro]

Is it me or the following two lines should be swapped?

annax/src/annax/ops.py

Lines 51 to 52 in 566aa3e

	j = jax.lax.cond(mask[i], lambda _: j + 1, lambda _: j, None)
	out = jax.lax.cond(mask[i], lambda _: out.at[j].set(i), lambda _: out, None)

[webbrain-one]

Good catch! 👀 I've opened PR #17 to swap those lines. Thanks for pointing this out!