jax

768b1d6d - Check for args_not_mapped first and then fancy_primitive_batchers. If not, it leads to numerical issues (no idea why, but this was the behavior before skippable_batchers was deleted)