Upgrade jaxlib to 0.4.31

[eightysteele]

This PR upgrades jaxlib to 0.4.31 with a few important caveats:

1. To get jax and torch working together peacefully with CUDA, I temporarily removed support for osx-arm64. This means that one will not be able to install b3d on a Mac to use the gcp tooling (but we have pretty good documentation for how to do this manually here. There’s a broader question here about if we should support osx-arm64 at all, which we can address separately within the context of data collected from the B3D user survey).
2. Tests are not currently passing on the latest commit of gen3d (378f93e as of this writing), so it’s hard to reason about if the upgraded jaxlib in this PR breaks anything. To help assess that, I captured the output of running tests on both branches, available here. Before merging this PR it's probably wise to study that a bit to see if there are any obvious regressions, etc.
3. Minor detail, but you no longer need to specify the GPU environment with -e gpu since it is now the default environment.