code: def test_poolnetworkbackprop(getkey):
      class CNN(eqx.Module):
          conv_layer: list[Union[eqx.nn.Conv2d, eqx.nn.MaxPool2d]]
          linear_layers: list[eqx.nn.Linear]
  
          def __init__(self, key):
              key1, key2, key3 = jax.random.split(key, 3)
              self.conv_layer = [
                  eqx.nn.Conv2d(3, 2, 3, key=key1),
                  eqx.nn.MaxPool2d(2, 2),
              ]
              self.linear_layers = [
                  eqx.nn.Linear(450, 256, key=key2),
                  eqx.nn.Linear(256, 10, key=key3),
              ]
  
          def __call__(self, x):
              for layer in self.conv_layer:
                  x = layer(x)
              x = jnp.ravel(x)
              for layer in self.linear_layers:
                  x = layer(x)
                  x = jax.nn.relu(x)
              return x
  
      cnn = CNN(getkey())
  
      @jax.vmap
      @jax.value_and_grad
      def loss_grad(x, y):
          return jax.numpy.mean((y - cnn(x)) ** 2)
  
      x = jrandom.normal(getkey(), (10, 3, 32, 32))
      y = jrandom.normal(getkey(), (10, 10))
      loss_grad(x, y)
  
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 46 for calling a global function with example usage: test_poolnetworkbackprop(getkey) and returns: jax, x