env_gpu.yaml

name: adept-gpu
channels:
  - conda-forge
dependencies:
  - python=3.12
  - ipython
  - ipykernel
  - pip
  - pip:
      # works for regular pip packages
      - --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
      - jaxlib==0.4.28+cuda12.cudnn89
      - jax==0.4.28
      - jaxopt
      - numpy
      - scipy
      - matplotlib
      - pyhdf
      - xlrd
      - pyyaml
      - mlflow
      - boto3
      - typing-extensions
      - optax
      - tqdm
      - xarray
      - mlflow_export_import
      - pint
      - diffrax
      - flatdict
      - plasmapy
      - parsl
      - interpax
      - tabulate