Project description

JaxHelper

Basic tools and helpers for Jax

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)

These details have not been verified by PyPI

View statistics for this project via Libraries.io, or by using our public dataset on Google BigQuery

This version

0.0.5

Jun 5, 2023

0.0.4

Jun 5, 2023

0.0.3

Jun 5, 2023

0.0.2

Jun 5, 2023

0.0.1

Jun 4, 2023

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

Uploaded Jun 5, 2023 Source

Uploaded Jun 5, 2023 Python 3

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

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