Code

[nnx] improve custom_vjp ergonomics#4489

Code

Open

cgarciae wants to merge 2 commits intonnx-optimize-fingerprintgoogle/flax:nnx-optimize-fingerprintfrom nnx-refactor-custom-vjpgoogle/flax:nnx-refactor-custom-vjpCopy head branch name to clipboard

+2,805-1,121

Commits

Commits on Jan 17, 2025

• [nnx] add cache_args

  cgarciae

  committed
  a9001c8

  Copy full SHA for a9001c8

Commits on Jan 19, 2025

• [nnx] improve custom_vjp ergonomics

  cgarciae

  committed
  1b7bab3

  Copy full SHA for 1b7bab3