code: def test_conv2d(getkey):
      # Positional arguments
      conv = eqx.nn.Conv2d(1, 3, 3, key=getkey())
      x = jrandom.normal(getkey(), (1, 32, 32))
      assert conv(x).shape == (3, 30, 30)
  
      # Some keyword arguments
      conv = eqx.nn.Conv2d(
          1, out_channels=3, kernel_size=(3, 3), key=getkey()
      )
      x = jrandom.normal(getkey(), (1, 32, 32))
      assert conv(x).shape == (3, 30, 30)
  
      # All keyword arguments
      conv = eqx.nn.Conv2d(
          in_channels=1,
          out_channels=3,
          kernel_size=(3, 3),
          padding=1,
          use_bias=False,
          key=getkey(),
      )
      x = jrandom.normal(getkey(), (1, 32, 32))
      assert conv(x).shape == (3, 32, 32)
  
      # Test strides
      conv = eqx.nn.Conv2d(
          in_channels=3,
          out_channels=1,
          kernel_size=(3, 3),
          stride=2,
          padding=1,
          use_bias=True,
          key=getkey(),
      )
      x = jrandom.normal(getkey(), (3, 32, 32))
      assert conv(x).shape == (1, 16, 16)
  
      # Test value matches
      conv = eqx.nn.Conv2d(1, 1, kernel_size=3, padding=1, key=getkey())
      new_weight = jnp.arange(9).reshape(1, 1, 3, 3)
      new_bias = jnp.array([1]).reshape(1, 1, 1)
      data = jnp.arange(-4, 5).reshape(1, 3, 3)
      assert new_weight.shape == conv.weight.shape
      assert new_bias.shape == conv.bias.shape  # pyright: ignore
      conv = eqx.tree_at(
          lambda x: (x.weight, x.bias), conv, (new_weight, new_bias)
      )
      answer = jnp.array(
          [-37, -31, -9, 25, 61, 49, 23, 41, 27]
      ).reshape(1, 3, 3)
      assert jnp.allclose(conv(data), answer)
  
      # Test groups
      conv = eqx.nn.Conv2d(
          2, 2, kernel_size=3, padding=1, key=getkey(), groups=2
      )
      # we will duplicate the weights from the "value matches" case
      # and multiply one copy by 2. Also, we modify the bias
      new_weight = jnp.concatenate(
          [
              1 * jnp.arange(9).reshape(1, 1, 3, 3),
              2 * jnp.arange(9).reshape(1, 1, 3, 3),
          ],
          axis=0,
      )
      new_bias = jnp.array([1, 2]).reshape(2, 1, 1)
  
      data = jnp.broadcast_to(
          jnp.arange(-4, 5).reshape(1, 3, 3),
          (2, 3, 3),
      )
      assert new_weight.shape == conv.weight.shape
      assert new_bias.shape == conv.bias.shape  # pyright: ignore
      conv = eqx.tree_at(
          lambda x: (x.weight, x.bias), conv, (new_weight, new_bias)
      )
      # this is the multiplication part, without the bias
      answer_part = jnp.array(
          [-38, -32, -10, 24, 60, 48, 22, 40, 26]
      ).reshape(1, 3, 3)
      answer = (
          jnp.concatenate(
              [
                  1 * answer_part,
                  2 * answer_part,
              ],
              axis=0,
          )
          + new_bias
      )
      assert jnp.allclose(conv(data), answer)
  
imports: import warnings
  from typing import Union
  import equinox
  import jax
  import jax.nn
  import jax.numpy
  import jax.random
  import pytest
label_desc: Use image node_id 14 for calling a global function with example usage: test_conv2d(getkey) without return types