Running tabulate on something with nn.vmap yields: 'ArrayImpl' object has no attribute '_state'

[mikesol]

Hi all!

I'm a newb and I'm using nn.vamp for the first time. When I run:

class Simple(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.vmap(
            lambda x: nn.Dense(features=8)(x),
            variable_axes={"params": 0},
            split_rngs={"params": True},
            in_axes=-2,
        )(x)
        return x

if __name__ == '__main__':
    ipt = 2**3
    print(Simple().tabulate(jax.random.key(0), jnp.ones((42, ipt)) ))

I'm getting:

(.venv) mikesol@Mikes-MacBook-Pro jaxfun % python model.py
Traceback (most recent call last):
  File "/Users/mikesol/devel/jaxfun/model.py", line 244, in <module>
    print(Simple().tabulate(jax.random.key(0), jnp.ones((42, ipt)) ))
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/mikesol/devel/jaxfun/.venv/lib/python3.11/site-packages/flax/linen/module.py", line 2464, in tabulate
    return tabulate_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/mikesol/devel/jaxfun/.venv/lib/python3.11/site-packages/flax/linen/summary.py", line 259, in _tabulate_fn
    table = table_fn(rngs, *fn_args, mutable=mutable, **fn_kwargs, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/mikesol/devel/jaxfun/.venv/lib/python3.11/site-packages/flax/linen/summary.py", line 279, in _get_table_fn
    variables = jax.eval_shape(_get_variables)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/mikesol/devel/jaxfun/.venv/lib/python3.11/site-packages/flax/linen/summary.py", line 277, in _get_variables
    return module.init(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/mikesol/devel/jaxfun/model.py", line 225, in __call__
    x = nn.vmap(
        ^^^^^^^^
AttributeError: 'ArrayImpl' object has no attribute '_state'

As this is my first time working with vmap and flax in general, I'm probably missing something obvious, but I can't quite figure it out from the error message. Any tips would be much appreciated! 🙏

[mikesol]

Digging in the docs, it seems like I may be using vmap wrong. I see that it accepts as the first argument:

• target – a Module or a function taking a Module as its first argument.

But my function above is a lambda that accepts input as its first argument. So a better question than the one above may be "how do I thread features=8 to the nn.Dense when using vmap?

[mikesol]

Answering my own question :)

class Simple(nn.Module):
    @nn.compact
    def __call__(self, x):
        DenseVMapped = nn.vmap(
            nn.Dense,
            variable_axes={"params": 0},
            split_rngs={"params": True},
            in_axes=-2,
        )
        x = DenseVMapped(features=42, name="foo")(x)
        return x

I couldn't find something like this in the docs - maybe updating a lifted example to show how to thread arguments would help? At any rate, sorry for the noise!