JaxHelper

Basic tools and helpers for Jax

Getting Started

python3 -m pip install jaxhelper

This repository contains basic helper functions I use every day.
Here are some highlights:

remat: function decorator to rematerialize ("activation checkpointing") hidden states during backward pass
softmax:
- exp in fp32 and matmul in bf16 (-> improved convergence and speed)
- fewer stored intermediates yet faster gradient
attention:
- recomputation of hidden states
- memory of O(K * N) rather than O(N^2) thanks to Self-attention Does Not Need O(n2) Memory. (No slowdown)

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

jaxhelper-0.0.5.tar.gz (3.6 kB view details)

Uploaded Jun 5, 2023 Source

If you're not sure about the file name format, learn more about wheel file names.

The dropdown lists show the available interpreters, ABIs, and platforms. Enable javascript to be able to filter the list of wheel files.

Uploaded Jun 5, 2023 Python 3

Details for the file jaxhelper-0.0.5.tar.gz.

Hashes for jaxhelper-0.0.5.tar.gz
Algorithm	Hash digest
SHA256	`226e7d4f6d5ad2c7910d51269180dd6243b57f73a22ce733a7daf1749eec5171`
MD5	`9cbff12828e4836743dd8c65a74d3670`
BLAKE2b-256	`759c8cbb45c6d5102abd0d85b7fe3fef2d364bf5c5c8715cbfdec7c3516d7707`

Details for the file jaxhelper-0.0.5-py3-none-any.whl.

Hashes for jaxhelper-0.0.5-py3-none-any.whl
Algorithm	Hash digest
SHA256	`8abd779dc88c627b73390f77bccd4cdf73dee08da0c6631c6c2cfe4d328484cf`
MD5	`cdf94f9d68f3a330cc632fa799501f55`
BLAKE2b-256	`c5d53a3251e3facaf123f50977fb1beeaa2254da9044e6af35ced3cfe50b96cc`