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)

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution for jaxhelper 0.0.5
File	Size	Uploaded
jaxhelper-0.0.5.tar.gz	3.6 kB	Jun 5, 2023	Details

Table of built distributions (wheels) for jaxhelper 0.0.5
File	Interpreter	ABI	Platform	Reset
jaxhelper-0.0.5-py3-none-any.whl 3.6 kB Jun 5, 2023	Python 3	none	any	Details

Total release size: 7.2 kB

Download URL	jaxhelper-0.0.5.tar.gz
Size	3.6 kB
Tags	Source
SHA-256 checksum How to use checksums	`226e7d4f6d5ad2c7910d51269180dd6243b57f73a22ce733a7daf1749eec5171`
BLAKE2b-256 checksum How to use checksums	`759c8cbb45c6d5102abd0d85b7fe3fef2d364bf5c5c8715cbfdec7c3516d7707`
Upload date	Jun 5, 2023
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/4.0.2 CPython/3.10.6`

Download URL	jaxhelper-0.0.5-py3-none-any.whl
Size	3.6 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`8abd779dc88c627b73390f77bccd4cdf73dee08da0c6631c6c2cfe4d328484cf`
BLAKE2b-256 checksum How to use checksums	`c5d53a3251e3facaf123f50977fb1beeaa2254da9044e6af35ced3cfe50b96cc`
Upload date	Jun 5, 2023
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/4.0.2 CPython/3.10.6`