Repository: google/flax

Issue:
Transpose Convolution module issue when used with circular padding and vmap
### Problem you have encountered:
I'm simply trying to `vmap` a `ConvTranspose` layer with circular padding, and it results in a `jax._src.errors.TracerArrayConversionError`. I'm running things on GPU.

### Steps to reproduce:
Here is a minimum example that reproduces the error.
```
import jax
import jax.numpy as jnp
import flax.linen as nn

layer = nn.ConvTranspose(features=5, kernel_size=(3,), padding="CIRCULAR")

# this is ok
sample_input = jnp.ones((1, 32, 2))
out, vars = layer.init_with_output(jax.random.PRNGKey(0), sample_input)
print(out.shape)

batch_input = jnp.ones((8, 4, 32, 2))
batch_apply = jax.vmap(layer.apply, in_axes=(None, 0))

# this breaks with the error provided
batch_out = batch_apply(vars, batch_input)
print(batch_out.shape)
```
I'm moderately confident that this is a bug that is specific to the transpose convolution because I verified that the code works ok if `nn.ConvTranspose` is replaced with `nn.Conv`. Things are also ok when `vmap` is not used.

### Logs, error messages, etc:

```
UnfilteredStackTrace: jax._src.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ShapedArray(float32[4,34,5])>with<BatchTrace(level=1/0)> with
  val = DeviceArray([[[[ 0.19632685,  0.56257343,  0.6327205 ,  0.278047  ,
...
```

### My guess at what's happening:
I'm suspecting that since padding needs to be added, the shape information becomes not static any more at some point?






Discussion:

Resolving patch:
diff --git a/flax/linen/linear.py b/flax/linen/linear.py
--- a/flax/linen/linear.py
+++ b/flax/linen/linear.py
@@ -603,7 +603,7 @@ def __call__(self, inputs: Array) -> Array:
       total_pad = [
           ((size_diff + 1) // 2, size_diff // 2) for size_diff in size_diffs
       ]
-      y = np.pad(y, [(0, 0)] + total_pad + [(0, 0)])
+      y = jnp.pad(y, [(0, 0)] + total_pad + [(0, 0)])
       # Wrap the result periodically around each spatial dimension,
       # one by one.
       for i in range(1, y.ndim - 1):


Test patch:
diff --git a/tests/linen/linen_linear_test.py b/tests/linen/linen_linear_test.py
--- a/tests/linen/linen_linear_test.py
+++ b/tests/linen/linen_linear_test.py
@@ -773,6 +773,21 @@ def test_circular_conv_transpose_2d_constant(
     )
     np.testing.assert_allclose(y, correct_ans)
 
+  def test_circular_conv_transpose_2d_with_vmap(self):
+    layer = nn.ConvTranspose(features=5, kernel_size=(3,), padding="CIRCULAR")
+
+    # this is ok
+    sample_input = jnp.ones((1, 32, 2))
+    out, vars = layer.init_with_output(jax.random.PRNGKey(0), sample_input)
+    self.assertEqual(out.shape, (1, 32, 5))
+
+    batch_input = jnp.ones((8, 32, 2))
+    batch_apply = jax.vmap(layer.apply, in_axes=(None, 0))
+
+    # this breaks with the error provided
+    batch_out = batch_apply(vars, batch_input)
+    self.assertEqual(batch_out.shape, (8, 32, 5))
+
   def test_circular_conv_transpose_1d_custom(self):
     """Test 1d transposed convolution with circular padding and a stride."""
     rng = dict(params=random.PRNGKey(0))