code: def _pad_norm_assert_shape(outputs):
      """Test that the spatial shape output is correct for masked BN/GN."""
      outputs, spatial_shape = outputs
      # Does it have the correct shape?
      np.testing.assert_equal(
          spatial_shape.shape, MaskedLayersTest.INPUTS_SHAPE_SMALL.shape
      )
      # Inferred spatial shape has the correct values.
      np.testing.assert_allclose(
          spatial_shape, MaskedLayersTest.INPUTS_SHAPE_SMALL, atol=0
      )
  
      _, h, w, _ = MaskedLayersTest.INPUTS_SMALL.shape
      norm_unpad = outputs[:, :h, :w, :]  # Region without padding.
      norm_right_pad = outputs[:, :h, w:, :]  # Right padded region.
      norm_bottom_pad = outputs[:, h:, :w]  # Bottom padded region.
      return norm_unpad, [norm_right_pad, norm_bottom_pad]
  
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 1 for calling a global function with example usage: _pad_norm_assert_shape(outputs) and returns: norm_unpad,  list