name: tfm_env

channels:
  - conda-forge
  - defaults
  - anaconda
dependencies:
  - jupyterlab
  - pip
  - python=3.10
  - pip:
    - huggingface_hub[cli]
    - utilsforecast
    - praxis
    - paxml
    - -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html  
    - "jax[cuda12_pip]"
    - einshape