Add support for Singularity/Apptainer container

[runame]

Resolves #436.

@pomonam Could you test this setup on your university's slurm cluster?

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅

[runame]

The changes to the Dockerfile seem to break the tests. Locally building a docker container from the modified Dockerfile worked just fine. @priyakasimbeg Any idea what's going on?

[priyakasimbeg]

The changes to the Dockerfile seem to break the tests. Locally building a docker container from the modified Dockerfile worked just fine. @priyakasimbeg Any idea what's going on?

@runame that is strange, these tests don't use Docker. It seems like the failure is coming from the flax import:

Traceback (most recent call last):
  File "/home/runner/work/algorithmic-efficiency/algorithmic-efficiency/tests/reference_algorithm_tests.py", line 32, in <module>
    import flax
  File "/opt/hostedtoolcache/Python/3.9.18/x64/lib/python3.9/site-packages/flax/__init__.py", line 18, in <module>
    from .configurations import (
  File "/opt/hostedtoolcache/Python/3.9.18/x64/lib/python3.9/site-packages/flax/configurations.py", line 25, in <module>
    from jax import config as jax_config
  File "/opt/hostedtoolcache/Python/3.9.18/x64/lib/python3.9/site-packages/jax/__init__.py", line 35, in <module>
    from jax import config as _config_module
  File "/opt/hostedtoolcache/Python/3.9.18/x64/lib/python3.9/site-packages/jax/config.py", line 17, in <module>
    from jax._src.config import config  # noqa: F401
  File "/opt/hostedtoolcache/Python/3.9.18/x64/lib/python3.9/site-packages/jax/_src/config.py", line 24, in <module>
    from jax._src import lib
  File "/opt/hostedtoolcache/Python/3.9.18/x64/lib/python3.9/site-packages/jax/_src/lib/__init__.py", line 92, in <module>
    import jaxlib.xla_client as xla_client
  File "/opt/hostedtoolcache/Python/3.9.18/x64/lib/python3.9/site-packages/jaxlib/xla_client.py", line 225, in <module>
    float8_e4m3b11fnuz = ml_dtypes.float8_e4m3b11
AttributeError: module 'ml_dtypes' has no attribute 'float8_e4m3b11'

These small tests are run on GH runners w CPU. I think maybe some "sub" package has not been pinned. I'm also seeing this in my other PRs. Let's not block this PR on those tests though.
We should just merge this once Juhan confirms the above procedure works for him too.

[priyakasimbeg]

Regarding failing tests; I filed jax-ml/jax#17693 w jax. I think we should just fix ml_dtypes to 0.2.0 in the meantime while they fix their dependencies.

Update: I pinned the ml_dtypes package in #514. If you sync this branch to dev, the tests should pass.

[runame]

Update: I pinned the ml_dtypes package in #514. If you sync this branch to dev, the tests should pass.

Cool, thanks! The tests pass now.