jax_examples

Deep Learning examples using the Jax ecosystem of libraries

env setup

poetry shell
poetry install --with dev

# inside the project root
jupyter lab

Alternatively, you can use vscode or the provided devcontainer with vscode to run the project in a containerised environment.

Name		Name	Last commit message	Last commit date
Latest commit History 243 Commits
.devcontainer		.devcontainer
.github		.github
checkpoints		checkpoints
images		images
playground		playground
utils		utils
.dockerignore		.dockerignore
.gitignore		.gitignore
01_basic_matrix.ipynb		01_basic_matrix.ipynb
02_differentiation.ipynb		02_differentiation.ipynb
03_diff_on_vectors.ipynb		03_diff_on_vectors.ipynb
04_random_numbers_and_arrays.ipynb		04_random_numbers_and_arrays.ipynb
05_jit.ipynb		05_jit.ipynb
06_vmap.ipynb		06_vmap.ipynb
07_autodiff_and_linear_regression.ipynb		07_autodiff_and_linear_regression.ipynb
08_ann_optax_flax.ipynb		08_ann_optax_flax.ipynb
09_lin_reg_optax.ipynb		09_lin_reg_optax.ipynb
10_mini_batches_dataloading.ipynb		10_mini_batches_dataloading.ipynb
11_pos_tag_lstm.ipynb		11_pos_tag_lstm.ipynb
12_cnn_cifar10.ipynb		12_cnn_cifar10.ipynb
13_model_checkpointing.ipynb		13_model_checkpointing.ipynb
14_fine_tune_bert.ipynb		14_fine_tune_bert.ipynb
LICENSE		LICENSE
README.md		README.md
azure-pipelines.yml		azure-pipelines.yml
ci.Dockerfile		ci.Dockerfile
devcontainer.Dockerfile		devcontainer.Dockerfile
env.yml		env.yml
poetry.lock		poetry.lock
pyproject.toml		pyproject.toml