Working Notes: a commonplace notebook for recording & exploring ideas.
Home. Site Map. Subscribe. More at expLog. — Kunal

Jax

Quick notes on jax

Tutorial
import jax.numpy as jnp
jax.vmap use this to auto-batch over dimensions; can specify axes
jax.shard_map to map over shards for manual parallelism