Solution to issue #2489

[jamesheald]

I believe this PR resolves issue #2489.

The solution is to perform code branching using the conditional control flow jax.lax.cond instead of jp.where.

From what I understand, jax.lax.cond only evaluates the branch that is used, whereas jp.where evaluates both the used and the unused branches. Consequently, if the unused branch has an NaN, this will cause a problem for jp.where (even though the branch is not executed) but not for jax.lax.cond.

@Balint-H noted that the NaNs encountered in issue #2489 go away if tausmooth is set to a nonzero value. I believe this is because the computation of tau_smooth involves dividing by smoothing_width, which gives NaN if smoothing_width = 0 (again, even if the smooth switching branch is not being executed).

If I perform branching with jax.lax.cond (as per this PR), the NaNs go away, even with tausmooth = 0.

To get the MWE provided in issue #2489 to work, I also had to replace jp.where with jax.lax.cond here, again because it involves dividing by 0. I have included this change in my PR. More generally, there are quite a lot of jp.where calls in the mjx codebase, and I wonder if these should also be changed to jax.lax.cond?

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.