class_docstr: Tests for modules in masked_layers.py.
class_docstr_tok: ["Tests","for","modules","in","masked_layers.py","."]
docstr: Pool on unpadded data and pool on padded and masked data are same.
docstr_tok: ["Pool","on","unpadded","data","and","pool","on","padded","and","masked","data","are","same","."]
code: def test_unpadded_pool_eq_masked_padded(
      self, pool_fn, window_shape, strides, padding
  ):
      """Pool on unpadded data and pool on padded and masked data are same."""
      pool_fn, masked_pool_fn = self.POOL_FN_DICT[pool_fn]
      output_small = pool_fn(
          self.INPUTS_SMALL, window_shape, strides, padding=padding
      )
      output_large = pool_fn(
          self.INPUTS_LARGE, window_shape, strides, padding=padding
      )
  
      outputs_padded, spatial_shape = masked_pool_fn(
          self.INPUTS_PADDED,
          window_shape,
          strides,
          padding=padding,
          spatial_shape=self.SPATIAL_SHAPE,
      )
  
      # Inferred spatial shape has the right shape.
      self.assertEqual(
          spatial_shape.shape, (self.INPUTS_PADDED.shape[0], 2)
      )
  
      # Inferred spatial shape has the right values.
      n_small, n_large = output_small.shape[0], output_large.shape[0]
      np.testing.assert_allclose(
          spatial_shape[:n_small, ...],
          np.stack(
              [np.array(output_small.shape[1:-1])] * n_small, axis=0
          ),
          atol=0,
      )
      np.testing.assert_allclose(
          spatial_shape[-n_large:, ...],
          np.stack(
              [np.array(output_large.shape[1:-1])] * n_large, axis=0
          ),
          atol=0,
      )
  
      # Masked output has the right values in the *un*masked region.
      ind_small = [slice(s) for s in output_small.shape[1:-1]]
      ind_large = [slice(s) for s in output_large.shape[1:-1]]
      ind_small = tuple([slice(n_small)] + ind_small + [slice(None)])
      ind_large = tuple(
          [slice(n_small, None)] + ind_large + [slice(None)]
      )
  
      np.testing.assert_allclose(
          output_small, outputs_padded[ind_small], atol=1e-5
      )
      np.testing.assert_allclose(
          output_large, outputs_padded[ind_large], atol=1e-5
      )
  
      # Masked output has the right values in the masked region.
      ind_small = [slice(s, None) for s in output_small.shape[1:-1]]
      ind_large = [slice(s, None) for s in output_large.shape[1:-1]]
      ind_small = tuple([slice(n_small)] + ind_small + [slice(None)])
      ind_large = tuple(
          [slice(n_small, None)] + ind_large + [slice(None)]
      )
  
      np.testing.assert_allclose(
          jnp.zeros_like(outputs_padded[ind_small]),
          outputs_padded[ind_small],
          atol=1e-5,
      )
      np.testing.assert_allclose(
          jnp.zeros_like(outputs_padded[ind_large]),
          outputs_padded[ind_large],
          atol=1e-5,
      )
  
imports: import dataclasses
  from typing import Any, Callable, Dict, Sequence, Tuple, Type, Union
  from absl.testing import absltest
  from absl.testing import parameterized
  import flax.linen
  from jax import random
  import jax.numpy
  import numpy
  from scenic.model_lib.layers import masked_layers
label_desc: Use image node_id 4 for calling the MaskedLayersTest obj's underlying member method code with example usage: obj.test_unpadded_pool_eq_masked_padded(pool_fn, window_shape, strides, padding) without return types