Dongseong Hwang

Before this PR, ParticleSystemTransformFeedback looks work but the logic is actually wrong against explanation in the book. Previously particles disappear during animation.

Introduce safe_not() for ~paddings instead of 1 - paddings.

1

Now all of paddings.dtype are bool (1bit), after https://github.com/apple/axlearn/pull/1133 This PR is unifying coding style between ~paddings and 1 - paddings used inconsistently across the codebase. I’d like to standardize...

`assert bool(jnp.isfinite(x).all()), f"Check numerics {msg_fmt.format(**msg_kwargs)}: {x}"` doesn't work with traced `x`. ``` assert bool(jnp.isfinite(x).all()), f"Check numerics {msg_fmt.format(**msg_kwargs)}: {x}" Traceback (most recent call last): File "/Users/dongseong/miniforge3/envs/ajax/lib/python3.10/site-packages/jax/_src/core.py", line 782, in __bool__ return...

Stop gradient on `residual = residual - z_q_i.detach()`

2

DAC doesn't stop gradient for `residual = residual - z_q_i.detach()`. https://github.com/descriptinc/descript-audio-codec/blob/c7cfc5d2647e26471dc394f95846a0830e7bec34/dac/nn/quantize.py#L186 However, vector_quantize_pytorch stop gradient on it. https://github.com/lucidrains/vector-quantize-pytorch/blob/976335f03b259536a06ed88a7adb7248cdb3d17c/vector_quantize_pytorch/residual_vq.py#L386 RVQ is trained greedily, stage by stage: Each quantizer should only be...