jax_examples

Deep Learning examples using the Jax ecosystem of libraries

env setup

python3 -m venv venv
source venv/bin/activate

pip install -r requirements.txt -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html --extra-index-url https://download.pytorch.org/whl/cpu

Name		Name	Last commit message	Last commit date
Latest commit History 193 Commits
.devcontainer		.devcontainer
.github		.github
checkpoints		checkpoints
images		images
playground		playground
.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
README.md		README.md
check_device.py		check_device.py
devcontainer.Dockerfile		devcontainer.Dockerfile
poetry.lock		poetry.lock
pyproject.toml		pyproject.toml
sanity_check.py		sanity_check.py

ShawonAshraf/jax_examples

Folders and files

Latest commit

History

Repository files navigation

jax_examples

env setup

About

Topics

Resources

Stars

Watchers

Forks

Languages