Optax recompiles when learning rate changes

[patrick-kidger]

Optax isn't treating its optimisers as pytrees. This is introducing spurious recompilation:

import functools as ft
import jax
import jax.numpy as jnp
import optax

optim1 = optax.adam(jnp.array(3e-3))
optim2 = optax.adam(jnp.array(3e-4))

@ft.partial(jax.jit, static_argnums=0)  # post-pytree-isation, this would probably just be `jax.jit`
def evaluate(optim, params):
    print("Compiling!")
    state = optim.init(params)
    grads = jnp.zeros_like(params)
    updates, new_state = optim.update(grads, state)
    new_params = optax.apply_updates(params, updates)
    return new_params, new_state

evaluate(optim1, jnp.array(0))  # Compiling!
evaluate(optim2, jnp.array(0))  # Compiling!

Can Optax optimisers be treated as pytrees with respect to their input arguments?

[patrick-kidger]

Okay, so for anyone else bumping into this issue, here's a closure_to_pytree function that solves the issue.

# Basically, Optax makes the decision *not* to register their optimisers as PyTrees.
# This means that we often end up with spurious recompilation, just because a learning
# rate changed. That results in a new optimiser instance, which is just a function and
# is treated statically.
#
# So here we simply replace all function closures with pytrees, with each of their cell
# contents as their subnodes.

import types
from typing import Any, Optional

import jax.tree_util as jtu

from equinox import Module  # https://github.com/patrick-kidger/equinox


def _make_cell(val):
    fn = lambda: val
    return fn.__closure__[0]  # pyright: ignore


def _adjust_function_closure(fn, closure):
    out = types.FunctionType(
        code=fn.__code__,
        globals=fn.__globals__,
        name=fn.__name__,
        argdefs=fn.__defaults__,
        closure=closure,
    )
    out.__module__ = fn.__module__
    out.__qualname__ = fn.__qualname__
    out.__doc__ = fn.__doc__
    out.__annotations__.update(fn.__annotations__)
    if fn.__kwdefaults__ is not None:
        out.__kwdefaults__ = fn.__kwdefaults__.copy()
    return out


# Not a pytree.
# Used so that two different local functions, with different identities, can still
# compare equal. This is needed as these leaves are compared statically when
# filter-jit'ing.
class _FunctionWithEquality:
    def __init__(self, fn: types.FunctionType):
        self.fn = fn

    def information(self):
        return self.fn.__qualname__, self.fn.__module__

    def __hash__(self):
        return hash(self.information())

    def __eq__(self, other):
        return type(self) == type(other) and self.information() == other.information()


class _Closure(Module):
    fn: _FunctionWithEquality
    contents: Optional[tuple[Any, ...]]

    def __init__(self, fn: types.FunctionType):
        self.fn = _FunctionWithEquality(fn)
        if fn.__closure__ is None:
            contents = None
        else:
            contents = tuple(
                closure_to_pytree(cell.cell_contents) for cell in fn.__closure__
            )
        self.contents = contents

    def __call__(self, *args, **kwargs):
        if self.contents is None:
            closure = None
        else:
            closure = tuple(_make_cell(contents) for contents in self.contents)
        fn = _adjust_function_closure(self.fn.fn, closure)
        return fn(*args, **kwargs)


def _fixup_closure(leaf):
    if isinstance(leaf, types.FunctionType):
        return _Closure(leaf)
    else:
        return leaf


def closure_to_pytree(tree):
    """Convert all function closures into pytree nodes.

    **Arguments:**

    - `tree`: Any pytree.

    **Returns:**

    A copy of `tree`, where all function closures have been replaced by a new object
    that is (a) callable like the original function, but (b) iterates over its
    `__closure__` as subnodes in the pytree.

    !!! Example

        ```python
        def some_fn():
            a = jnp.array(1.)

            @closure_to_pytree
            def f(x):
                return x + a

            print(jax.tree_util.tree_leaves(f))  # prints out `a`
        ```

    !!! Warning

        One annoying technical detail in the above example: we had to wrap the whole lot
        in a `some_fn`, so that we're in a local scope. Python treats functions at the
        global scope differently, and this conversion won't result in any global
        variable being treated as part of the pytree.

        In practice, the intended use case of this function is to fix Optax, which
        always uses local functions.
    """
    return jtu.tree_map(_fixup_closure, tree)

Example usage is:

lr = jnp.array(1e-3)
optim = optax.chain(
    optax.adam(lr),
    optax.scale_by_schedule(optax.piecewise_constant_schedule(1, {200: 0.1})),
)
optim = closure_to_pytree(optim)

[patrick-kidger]

(CC @cgarciae as I recall us talking about this way back when.)

[fabianp]

Closing issue given the snippet that @patrick-kidger shared solves the issue

[peabody124]

Have you guys considered incorporating this change upstream? This retracing behavior is super counterintuitive!

[colehaus]

I also encountered this recently. Note that it means that, for example, you will never hit the persistent compilation cache on your training steps if you're doing a hyperparameter sweep.

See the following example:

def main(lr: ndarray[bfloat16]):
    tx = optax.adafactor(learning_rate=lr, weight_decay_rate=lr, momentum=lr, dtype_momentum=bfloat16)
    model = jnp.ones(8, bfloat16)
    opt_state = tx.init(model)
    print(jax.jit(tx.update).lower(model, opt_state, model).as_text())

where the learning rate gets embedded into the HLO:

module @jit_update_fn attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<8xbf16>, %arg1: tensor<i32>, %arg2: tensor<8xbf16>, %arg3: tensor<i32>, %arg4: tensor<8xbf16>, %arg5: tensor<8xbf16>) -> (tensor<8xbf16> {jax.result_info = "[0]"}, tensor<i32> {jax.result_info = "[1][0].count"}, tensor<1xbf16> {jax.result_info = "[1][0].v_row"}, tensor<1xbf16> {jax.result_info = "[1][0].v_col"}, tensor<8xbf16> {jax.result_info = "[1][0].v"}, tensor<i32> {jax.result_info = "[1][4].count"}, tensor<8xbf16> {jax.result_info = "[1][4].ema"}) {
    %cst = stablehlo.constant dense<4.394530e-01> : tensor<bf16>
    %cst_0 = stablehlo.constant dense<4.394530e-01> : tensor<bf16>
    ...

Note that it can be worked around with inject_hyperparams.