fix: memoize optimizer factory functions to avoid JIT recompilation (#353)

[Muneerali199]

Fixes #353

Problem

When GradientTransformation objects are passed as static arguments to jax.jit, JAX recompiles on every call because each call to an optimizer factory creates new closures with different identities, producing different hashes.

Root Cause

GradientTransformation is a NamedTuple whose __hash__ is derived from the identity of its init and update closure fields. Since every call to optax.adam(1e-3) creates new closures, the hash is different each time.

Solution

Add @functools.lru_cache(maxsize=None) to all optimizer alias factory functions. This ensures that identical arguments always return the exact same GradientTransformation object with a stable hash.

Files changed:

• optax/_src/alias.py: Added memoization to all 29 optimizer factory functions
• optax/transforms/_combining.py: Added to chain()
• optax/_src/base.py: Added to identity(), set_to_zero(), stateless(), stateless_with_tree_map(), with_extra_args_support()
• optax/_src/alias_test.py: Added GradientTransformationMemoizationTest with 6 tests

Verification

• All memoized factory functions return is-identical objects for identical arguments
• Different arguments still produce different objects
• JAX jax.jit does not recompile when the same optimizer configuration is reused

[rdyro]

Thanks, this looks like an interesting direction!

The straightforward application lru_cache is probably going to break on any dynamic data which is not hashable, what are you thinking as far as solving that problem?

[Muneerali199]

Thanks, this looks like an interesting direction!

The straightforward application lru_cache is probably going to break on any dynamic data which is not hashable, what are you thinking as far as solving that problem?

Good point. I'd handle it by wrapping lru_cache to catch TypeError from unhashable args and fall through cleanly — no crash, no id() reuse bug. For the common case (hashable args) it's cached and stable. Want me to update the PR with this?