class_docstr: Tests for utilities in nn_ops.py.
class_docstr_tok: ["Tests","for","utilities","in","nn_ops.py","."]
docstr: Test Pooling module.
  
  Args:
    pooling_type: str; Type of pooling function from `['avg_pooling',
      'max_pooling', 'space_to_depth']`
docstr_tok: ["Test","Pooling","module",".","Args",":","pooling_type",":","str",";","Type","of","pooling","function","from","`","[","'avg_pooling","'",",","'max_pooling","'",",","'space_to_depth","'","]","`"]
code: def test_pooling(self, pooling_type):
      """Test Pooling module.
  
      Args:
        pooling_type: str; Type of pooling function from `['avg_pooling',
          'max_pooling', 'space_to_depth']`
      """
      inputs_shape = (16, 32, 32, 64)
      window_shape = (4, 4)
      strides = (4, 4)
      inputs = jnp.array(np.random.normal(size=inputs_shape))
  
      outputs = nn_ops.pooling(
          inputs,
          pooling_configs={"pooling_type": pooling_type},
          window_shape=window_shape,
          strides=strides,
      )
  
      if pooling_type == "space_to_depth":
          self.assertEqual(outputs.shape, (16, 8, 8, 1024))
      else:
          self.assertEqual(outputs.shape, (16, 8, 8, 64))
  
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  import flax.linen
  import jax
  import jax.numpy
  import numpy
  from scenic.model_lib.layers import nn_ops
label_desc: Use image node_id 8 for calling the NNOpsTest obj's underlying member method code with example usage: obj.test_pooling(pooling_type) without return types