Unable to install brax with Bazel via pip due to orbax-checkpoint

[hartikainen]

I'm using Bazel and installing python dependencies via pip using rules_python. I tried upgrading brax from 0.11.0 to 0.12.1 but couldn't get that to work because orbax-checkpoint, which is a dependency of brax, fails to install on Bazel. This is clearly an issue within orbax and not brax, but I wanted to ask here if anyone else has seen such issue or knows any workarounds for it. I've already raised this issue (google/orbax#1429) on their end but haven't heard back anything.