How to use nn.DenseGeneral

[Peter-Vincent]

I have the following network

def __call__(self,x):
        x = nn.DenseGeneral(features = x.shape[1:],axis=(1,2,3),batch_dims=(0,))(x)
        x = nn.Conv(features=64, kernel_size=(4,4))(x)
        x = nn.relu(x)
        x = nn.avg_pool(x, window_shape=(2,2), strides=(2,2))
        x = nn.Conv(features=128, kernel_size=(3,3))(x)
        x = nn.relu(x)
        x = nn.avg_pool(x, window_shape=(2,2), strides=(2,2))
        x = x.reshape((x.shape[0],-1))
        x = nn.Dense(features = 2000)(x)
        x = nn.relu(x)
        x = nn.Dense(features = 502)(x)
        x = nn.relu(x)
        return jnp.squeeze(x)

But I just don't understand how nn.DenseGeneral is working - everything I have tried so far errors out.

The images I want to pass in are 1 x 51 x 51 x 1. I want the first layer to be fully connected. The batch I am passing in is a numpy array with shape 32 x 51 x 51 x 1, so there are 32 images in each batch. I initialise the network with the default shape for the image, so
params = module.init(rng, jnp.ones([1, 51, 51, 1]))['params']
but everytime I try and run my network it complains with this
flax.errors.ScopeParamShapeError: Inconsistent shapes between value and initializer for parameter "kernel" in "/DenseGeneral_0": (1, 51, 51, 1, 51, 51, 1), (32, 51, 51, 1, 51, 51, 1). (https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html#flax.errors.ScopeParamShapeError)
I guess this is coming from it not really batching over that first dimension? At a push I can just use vmap to batch like that, but that will be very memory-expensive, so I'd rather not. Can someone suggest how to implement this network?

[cgarciae]

Hey @Peter-Vincent, to fix this just don't specify the batch_dims argument. Here is an example:

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

class Foo(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.DenseGeneral(features = x.shape[1:],axis=(1,2,3))(x)
        x = nn.Conv(features=64, kernel_size=(4,4))(x)
        x = nn.relu(x)
        x = nn.avg_pool(x, window_shape=(2,2), strides=(2,2))
        x = nn.Conv(features=128, kernel_size=(3,3))(x)
        x = nn.relu(x)
        x = nn.avg_pool(x, window_shape=(2,2), strides=(2,2))
        x = x.reshape((x.shape[0],-1))
        x = nn.Dense(features = 2000)(x)
        x = nn.relu(x)
        x = nn.Dense(features = 502)(x)
        x = nn.relu(x)
        return jnp.squeeze(x)

x = jax.random.normal(jax.random.PRNGKey(0), (1, 51, 51, 1))    
module = Foo()

print(module.tabulate(jax.random.PRNGKey(1), x))
variables = module.init(jax.random.PRNGKey(1), x)

x = jax.random.normal(jax.random.PRNGKey(0), (32, 51, 51, 1))
y = module.apply(variables, x)

print(jax.tree_map(lambda x: x.shape, variables))
print(y.shape)