DenseGeneral with more than 2 dimensions cannot be partitioned

DenseGeneral first computes a 2-dimensional "flat shape" to initialize its kernel (see [this line](https://github.com/google/flax/blob/d9585e0a6ba6c4a4ebc93d0707add573420703df/flax/linen/linear.py#L143)) and then later reshapes the kernel to the correct shape. However, the partitioning API (e.g., `nn.with_partitioning`) works by wrapping the kernel initializer. So if you have a DenseGeneral layer with a 3D kernel, and you try to partition this kernel, e.g.,

```Python
dense = partial(
nn.DenseGeneral,
axis=-1,
dtype=self.dtype,
features=(self.num_heads, head_dim),
kernel_init=nn.with_logical_partitioning(
self.kernel_init, ("embed", "heads", "head_dim")
),
)
```

JAX throws an error complaining about a 3D sharding constraint on a 2D array, because Flax tries to apply the constraint *before* the reshape (see [this line](https://github.com/google/flax/blob/d9585e0a6ba6c4a4ebc93d0707add573420703df/flax/linen/linear.py#L151)).