We install jax/flax/paxml/etc. as source installations in the containers. The git working copy that we check out in /opt has the same top-level folder name as the package themselves. This would cause Python to incorrectly import the package at the git repo level, if the Python script is run out of /opt.
Possible solutions:
- Rename git working copies to something like:
/opt/jax -> /opt/jax-source.
- Document and ask users to always use
/root or /workspace as the work area.