Deprecation of `nnx.Remat` and usage of `nnx.remat`

[lraedel]

Hi team,

I noticed that nnx.Remat has been deprecated but it is not easily visible from the changelog so it was a bit of a surprise. Trying to replace our usage of nnx.Remat with nnx.remat lead to several issues so I wanted to understand what the recommended usage pattern is.

The usage pattern we had and wanted to achieve is that if we have a layer such as the one below where we can dynamically add remat with different policies depending on the use-case.

import flax.nnx as nnx
import jax
import jax.numpy as jnp
class Linear(nnx.Module):
    def __init__(self, din: int, dout: int, rngs: nnx.Rngs):
        self.linear = nnx.Linear(din, dout, rngs=rngs)

    def __call__(self, x: jax.Array) -> jax.Array:
        return self.linear(x)

The thing we tried was to dynamically change __call__ which did not work and is most likely not a good idea.

Would appreciate if you could share the recommended usage pattern and if there is a way to achieve to re-use a module definition and dynamically apply remat to it with multiple policies.

Thank you!

[cgarciae]

Hey @lraedel ! If its useful maybe you could create a mini Remat abstraction:

import flax.nnx as nnx
import jax.numpy as jnp

def Remat(module, /, **remat_kwargs):
  class _Remat(nnx.Module):
    def __init__(self, module):
      self.module = module

    @nnx.remat(**remat_kwargs)
    def __call__(self, *args):
      return self.module(*args)

  return _Remat(module)

remat_linear = Remat(nnx.Linear(2, 3, rngs=nnx.Rngs(0)))
y = remat_linear(jnp.ones((1,2)))

[lraedel]

Thanks a lot and sorry for taking time to reply! Unfortunately, this does not work as I expected. When trying to call the remat module that wraps the underlying module it results in the following error when trying to call it with kwargs. I would have expected that the wrapped module behaves the same way in terms of how it can be called?

TypeError                                 Traceback (most recent call last)
Cell In[3], line 31
     29 y2 = remat_custom(x)
     30 y2 = remat_custom(x, 2.0)
---> 31 y2 = remat_custom(x, scale=2.0)

File /site-packages/flax/nnx/transforms/transforms.py:72, in resolve_kwargs.<locals>.resolve_kwargs_decorator.<locals>.resolve_kwargs_wrapper(*args, **kwargs)
     70 @functools.wraps(f)
     71 def resolve_kwargs_wrapper(*args, **kwargs):
---> 72   args = resolve_kwargs(f, args, kwargs)
     73   return f(*args)

File site-packages/flax/nnx/transforms/transforms.py:90, in resolve_kwargs(fun, args, kwargs)
     88 ba.apply_defaults()
     89 if ba.kwargs:
---> 90   raise TypeError('keyword arguments could not be resolved to positions')
     91 else:
     92   return ba.args

TypeError: keyword arguments could not be resolved to positions

The code to reproduce

from flax import nnx
import jax.numpy as jnp

def Remat(module, /, **remat_kwargs):
    class _Remat(nnx.Module):
        def __init__(self, module):
            self.module = module

        @nnx.remat(**remat_kwargs)
        def __call__(self, *args, **kwargs):
            return self.module(*args, **kwargs)

    return _Remat(module)

rng = nnx.Rngs(0)
linear = nnx.Linear(2, 3, rngs=rng)

class CustomModule(nnx.Module):
    def __init__(self, module):
        self.module = module
    
    def __call__(self, x, scale=1.0):
        return self.module(x) * scale

custom_module = CustomModule(linear)
remat_custom = Remat(custom_module)

x = jnp.ones((1, 2))
# works
y2 = remat_custom(x)
# works
y2 = remat_custom(x, 2.0)
# fails
y2 = remat_custom(x, scale=2.0)

Versions of flax and jax

jax                                      0.5.3
jaxlib                                   0.5.3
flax                                     0.10.5