Error during initialization when using flax.linen.vmap on a module with multiple called methods

[lucaslingle]

System information

• OS Platform and Distribution: MacOS Catalina
• Flax, jax, jaxlib versions: flax==0.6.8, jax==0.4.8, jaxlib==0.4.7
• Python version: 3.10
• GPU/TPU model and memory: N/A
• CUDA version (if applicable): N/A

Problem you have encountered:

Suppose SmallModule is a Linen module with multiple methods defined, including __call__. Suppose a vmapped version of SmallModule is declared in the setup method of another Linen module, LargeModule.

If LargeModule defines its __call__ method in a manner that first calls the vmapped SmallModule's __call__ method and then calls another method of the vmapped SmallModule, an error is raised.

The error says that certain parameters defined in SmallModule's setup method, are not attributes of VmapSmallModule. However, no error is not raised when calling both methods of SmallModule in a version of LargeModule NOT using vmap.

What you expected to happen:

The vmapped version of SmallModule should consistently behave as if it has the defined parameters.

Logs, error messages, etc:

Traceback (most recent call last):
  File "/Users/lucaslingle/PycharmProjects/empty_project/scripts/vmap_issue.py", line 97, in <module>
    main()
  File "/Users/lucaslingle/PycharmProjects/empty_project/scripts/vmap_issue.py", line 93, in main
    large_params = make_large_module(input_dim=3, hidden_dim=10, num_small=2)
  File "/Users/lucaslingle/PycharmProjects/empty_project/scripts/vmap_issue.py", line 83, in make_large_module
    params = LargeModule(input_dim, hidden_dim, num_small).init(
  File "/Users/lucaslingle/opt/miniconda3/envs/empty_project/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/lucaslingle/opt/miniconda3/envs/empty_project/lib/python3.10/site-packages/flax/linen/module.py", line 1621, in init
    _, v_out = self.init_with_output(
  File "/Users/lucaslingle/opt/miniconda3/envs/empty_project/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/lucaslingle/opt/miniconda3/envs/empty_project/lib/python3.10/site-packages/flax/linen/module.py", line 1526, in init_with_output
    return init_with_output(
  File "/Users/lucaslingle/opt/miniconda3/envs/empty_project/lib/python3.10/site-packages/flax/core/scope.py", line 965, in wrapper
    return apply(fn, mutable=mutable, flags=init_flags)({}, *args, rngs=rngs,
  File "/Users/lucaslingle/opt/miniconda3/envs/empty_project/lib/python3.10/site-packages/flax/core/scope.py", line 933, in wrapper
    y = fn(root, *args, **kwargs)
  File "/Users/lucaslingle/opt/miniconda3/envs/empty_project/lib/python3.10/site-packages/flax/linen/module.py", line 2102, in scope_fn
    return fn(module.clone(parent=scope), *args, **kwargs)
  File "/Users/lucaslingle/opt/miniconda3/envs/empty_project/lib/python3.10/site-packages/flax/linen/module.py", line 427, in wrapped_module_method
    return self._call_wrapped_method(fun, args, kwargs)
  File "/Users/lucaslingle/opt/miniconda3/envs/empty_project/lib/python3.10/site-packages/flax/linen/module.py", line 845, in _call_wrapped_method
    y = fun(self, *args, **kwargs)
  File "/Users/lucaslingle/PycharmProjects/empty_project/scripts/vmap_issue.py", line 57, in __call__
    output = self.smalls.call_transpose(hidden)
  File "/Users/lucaslingle/opt/miniconda3/envs/empty_project/lib/python3.10/site-packages/flax/linen/module.py", line 427, in wrapped_module_method
    return self._call_wrapped_method(fun, args, kwargs)
  File "/Users/lucaslingle/opt/miniconda3/envs/empty_project/lib/python3.10/site-packages/flax/linen/module.py", line 845, in _call_wrapped_method
    y = fun(self, *args, **kwargs)
  File "/Users/lucaslingle/PycharmProjects/empty_project/scripts/vmap_issue.py", line 22, in call_transpose
    return jnp.dot(x, self.proj.T)
  File "/Users/lucaslingle/opt/miniconda3/envs/empty_project/lib/python3.10/site-packages/flax/linen/module.py", line 922, in __getattr__
    raise AttributeError(msg)
jax._src.traceback_util.UnfilteredStackTrace: AttributeError: "VmapSmallModule" object has no attribute "proj".

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/lucaslingle/PycharmProjects/empty_project/scripts/vmap_issue.py", line 97, in <module>
    main()
  File "/Users/lucaslingle/PycharmProjects/empty_project/scripts/vmap_issue.py", line 93, in main
    large_params = make_large_module(input_dim=3, hidden_dim=10, num_small=2)
  File "/Users/lucaslingle/PycharmProjects/empty_project/scripts/vmap_issue.py", line 83, in make_large_module
    params = LargeModule(input_dim, hidden_dim, num_small).init(
  File "/Users/lucaslingle/PycharmProjects/empty_project/scripts/vmap_issue.py", line 57, in __call__
    output = self.smalls.call_transpose(hidden)
  File "/Users/lucaslingle/PycharmProjects/empty_project/scripts/vmap_issue.py", line 22, in call_transpose
    return jnp.dot(x, self.proj.T)
AttributeError: "VmapSmallModule" object has no attribute "proj".

Steps to reproduce:

import flax.linen as nn
import jax
import jax.numpy as jnp


class SmallModule(nn.Module):
    input_dim: int
    output_dim: int

    def setup(self):
        proj_args = [
            jax.nn.initializers.normal(1.0),
            [self.input_dim, self.output_dim],
            jnp.float32,
        ]
        self.proj = self.param("proj", *proj_args)

    def __call__(self, x):
        return jnp.dot(x, self.proj)

    def call_transpose(self, x):
        return jnp.dot(x, self.proj.T)


class LargeModuleNoVmap(nn.Module):
    input_dim: int
    hidden_dim: int

    def setup(self):
        self.small = SmallModule(
            input_dim=self.input_dim, output_dim=self.hidden_dim
        )

    def __call__(self, x):
        hidden = self.small(x)
        output = self.small.call_transpose(hidden)
        return hidden, output


class LargeModule(nn.Module):
    input_dim: int
    hidden_dim: int
    num_small: int

    def setup(self):
        self.smalls = nn.vmap(
            SmallModule,
            in_axes=0,
            out_axes=-2,
            variable_axes={"params": 0},
            split_rngs={"params": True},
        )(input_dim=self.input_dim, output_dim=self.hidden_dim // self.num_small)

    def __call__(self, x):
        tiled_x = jnp.tile(x[None, ...], [self.num_small] + [1] * jnp.ndim(x))
        hidden = self.smalls(tiled_x)
        output = self.smalls.call_transpose(hidden)
        hidden = jnp.reshape(hidden, list(hidden.shape[0:-2]) + [self.hidden_dim])
        output = jnp.reshape(output, list(output.shape[0:-2]) + [self.input_dim])
        return hidden, output


def make_small_module(input_dim, output_dim):
    xs = jnp.zeros([1, input_dim])
    params = SmallModule(input_dim, output_dim).init(
        {"params": jax.random.PRNGKey(0)},
        x=xs,
    )["params"].unfreeze()
    return params


def make_large_module_novmap(input_dim, hidden_dim):
    xs = jnp.zeros([1, input_dim])
    params = LargeModuleNoVmap(input_dim, hidden_dim).init(
        {"params": jax.random.PRNGKey(0)},
        x=xs,
    )["params"].unfreeze()
    return params


def make_large_module(input_dim, hidden_dim, num_small):
    xs = jnp.zeros([1, input_dim])
    params = LargeModule(input_dim, hidden_dim, num_small).init(
        {"params": jax.random.PRNGKey(0)},
        x=xs,
    )["params"].unfreeze()
    return params


def main():
    small_params = make_small_module(input_dim=3, output_dim=10)
    large_params = make_large_module_novmap(input_dim=3, hidden_dim=10)
    large_params = make_large_module(input_dim=3, hidden_dim=10, num_small=2)


if __name__ == "__main__":
    main()

[cgarciae]

Hey @lucaslingle, first thing you can do is use the methods argument for nn.vmap to specify you also want call_transpose to be lifted (by default only __call__ is lifted):

self.smalls = nn.vmap(
    SmallModule,
    in_axes=0,
    out_axes=-2,
    variable_axes={"params": 0},
    split_rngs={"params": True},
    methods=["__call__", "call_transpose"],
)(input_dim=self.input_dim, output_dim=self.hidden_dim // self.num_small)

This fixes the original error but now you have some logical errors. Hope this helps!

[lucaslingle]

Thanks, apologies for the silly question!