Equinox (JAX) neural nets.
Clone the repo:
git clone git@github.com:eringrant/nets.git
cd nets/
To install a Conda environment with the requisite packages on CPU:
conda env create --file environment-cpu.yml
To test the code a quick debug run:
python -m nets.experiments.in_context_learning.launcher_local
Optionally, define a few environment variables
by adding the following to a shell configuration file such as
~/.bashrc
, ~/.bash_profile
, ~/.bash_login
, or ~/.profile
:
export SCRATCH_HOME="..."
Then, follow one of two methods below to install nets
.
Use Conda to develop nets
code directly.
To install via Mamba (recommended) or Conda, do:
conda env create --file environment-cpu.yml
To install via Mamba (recommended) or Conda with GPU support, do:
conda env create --file environment-gpu.yml
If working on the head node of a SLURM cluster, you will need to create a GPU-compatible environment on a compute node with an available GPU via:
srun --partition=gpu --gres=gpu:1 conda env create -f environment-gpu.yml
Note that you may have to adapt the partition
name to the available partitions on your cluster;
run sinfo -s
to display details about partitions.
Use Pip in order to install this code as a package in another project.
python -m pip install git+https://github.com/eringrant/nets
nets @ git+https://github.com/eringrant/nets
git clone git@github.com:eringrant/nets.git
cd nets/
python -m pip install -e .
.pre-commit-config.yaml
has been configured to run several autoformatters.
Run the following to install, update, and cache all pre-commit tools:
pre-commit install && pre-commit run