Skip to content

Instantly share code, notes, and snippets.

@yberreby

yberreby/jax_lessons_learned.md

Created June 23, 2025 22:06
Show Gist options
  • Save yberreby/79850710c020dbe8d36828c54cde8e82 to your computer and use it in GitHub Desktop.
Save yberreby/79850710c020dbe8d36828c54cde8e82 to your computer and use it in GitHub Desktop.
Download ZIP
JAX Lessons Learned
Raw
jax_lessons_learned.md
  • Use jax.tree_util.Partial to pass partially-applied functions to JIT-compiled code
  • Use static_argnames, not static_argnums, whenever possible
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment