semvec/pypi-packages huggingface dataset:
tax-dp 0.1.1
A fully sharded data parallel trainer for jax
License: MIT
Requires Python: >=3.10

# jax-trainer
An efficient distributed FSDP in jax

# Running Examples
## Single host trainer
Run imdb movie classification from the long range arena dataset with batch parallelisation given by jit function:
```
pdm run bash ./tax/examples/bin/run_lra.sh $PWD imdb.yaml lra_imdb
```

For Library users see example usages in:

```
tax/tax/examples/lra/exp.py
```
## FSDP Trainer
Train a gemma language model with the FSDP trainer
```
pdm run bash ./tax/examples/bin/run_lm.sh $PWD lm.yaml test_trainer
```

For Library users see example usages in:

```
tax/tax/examples/gemma/exp.py
```