jax

JAX: composable transformations of Python+ NumPy programs(Bradbury et al. 2018)

https://github.com/google/jax

JAX in Action(“Welcome \(\cdot\) JAX in Action MEAP V02,” n.d.) * Compared to pytorch and tensorflow which organizes code in classes, jax typically organizes code into functions. * Jax is minimalistic, no data loaders no high level model components, those are being backfilled by other packages.

JAX Quickstart

Get started with JAX!

Jax Vs PyTorch Key Differences

JAX Crash Course - Accelerating Machine Learning code!

Why You Should (or Shouldn’t) be Using Google’s JAX in 2022

The Autodiff Cookbook

References

Bradbury, James, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, and Skye Wanderman-Milne. 2018. “JAX: Composable Transformations of Python+ NumPy Programs.” Version 0.2 5: 14–24.

“Welcome \(\cdot\) JAX in Action MEAP V02.” n.d. https://livebook.manning.com/book/jax-in-action/welcome/v-2/.