Depthwise convolution in `flax.nnx.Conv` is significantly slower than PyTorch and TensorFlow

### Description

I have noticed that the depthwise convolution using `nnx.Conv` is much slower compared to PyTorch and TensorFlow. I ran a quick benchmark to compare performance across the three frameworks, and the results show a huge gap in speed for the JAX version.

Here's a minimal example I used for benchmarking:

```python
import flax.nnx as nnx
import torch
import numpy as np
import jax
import jax.numpy as jnp
import tensorflow as tf

dim = 256
dwconv_pt = torch.nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim)
dwconv_jax = nnx.Conv(
in_features=dim,
out_features=dim,
kernel_size=3,
strides=1,
padding=1,
use_bias=True,
feature_group_count=dim,
input_dilation=1,
kernel_dilation=1,
rngs= nnx.Rngs(0)
)
dwconv_tf = tf.keras.layers.DepthwiseConv2D(
kernel_size=(3, 3),
strides=(1, 1),
padding='same',
depth_multiplier=1,
use_bias=True
)

np_input = np.random.rand(1, dim, 124, 124).astype('float32')
jax_input = jnp.array(np.transpose(np_input, (0, 2, 3, 1)))
tf_input = tf.convert_to_tensor(np.transpose(np_input, (0, 2, 3, 1)))
torch_input = torch.from_numpy(np_input)

# Benchmarking
import time
for _ in range(10):
dwconv_pt(torch_input)
dwconv_jax(jax_input)
dwconv_tf(tf_input)

# Timings
start_time = time.time()
for _ in range(50):
dwconv_pt(torch_input)
pt_time = time.time() - start_time

start_time = time.time()
for _ in range(50):
dwconv_jax(jax_input)
jax_time = time.time() - start_time

start_time = time.time()
for _ in range(50):
dwconv_tf(tf_input)
tf_time = time.time() - start_time

print(f'Tensorflow time: {tf_time}s')
print(f'PyTorch time: {pt_time}s')
print(f'JAX time: {jax_time}s')
```

## Results:
```python
TensorFlow time: 0.77s
PyTorch time: 0.49s
JAX (Flax) time: 36.57s
```

### System info (python version, jaxlib version, accelerator, etc.)

```
Python version: 3.10.12
JAX version: 0.4.26
Flax version: 0.8.4

Architecture: x86_64
CPU op-mode(s): 32-bit, 64-bit
Address sizes: 46 bits physical, 48 bits virtual
Byte Order: Little Endian
CPU(s): 2
On-line CPU(s) list: 0,1
Vendor ID: GenuineIntel
Model name: Intel(R) Xeon(R) CPU @ 2.20GHz
CPU family: 6
Model: 79
Thread(s) per core: 2
Core(s) per socket: 1
Socket(s): 1
Stepping: 0
BogoMIPS: 4399.99
Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 cl
flush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc re
p_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3
fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand
hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp
fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm rdseed adx sm
ap xsaveopt arat md_clear arch_capabilities
```