class_docstr: Tests Dataset Utilities.
class_docstr_tok: ["Tests","Dataset","Utilities","."]
docstr: Tests maybe_pad_batch.
docstr_tok: ["Tests","maybe_pad_batch","."]
code: def test_maybe_pad_batch(
      self,
      pixel_level,
      batch_mask_shape,
      mask_sum_coef,
      pre_padding_mask,
  ):
      """Tests maybe_pad_batch."""
      desired_bs = 32
      partial_bs = 13
  
      def make_fake_batches():
          # Make dummy batches
          complete_batch = {
              "inputs": jnp.array(
                  np.random.normal(size=(desired_bs, 28, 28, 3))
              ),
              "label": jnp.array(
                  np.random.normal(size=(desired_bs, 10))
              ),
          }
          partial_batch = {
              "inputs": jnp.array(
                  np.random.normal(size=(partial_bs, 28, 28, 3))
              ),
              "label": jnp.array(
                  np.random.normal(size=(partial_bs, 10))
              ),
          }
  
          complete_batch_mask, partial_batch_mask = None, None
          if pre_padding_mask:
              complete_batch_mask = jnp.broadcast_to(
                  jnp.eye(28, 28), (desired_bs, 28, 28)
              )
              complete_batch["batch_mask"] = complete_batch_mask
              partial_batch_mask = jnp.broadcast_to(
                  jnp.eye(28, 28), (partial_bs, 28, 28)
              )
              partial_batch["batch_mask"] = partial_batch_mask
  
          return (
              complete_batch,
              partial_batch,
              complete_batch_mask,
              partial_batch_mask,
          )
  
      ######## Test complete training batches
      (
          complete_batch,
          partial_batch,
          complete_batch_mask,
          partial_batch_mask,
      ) = make_fake_batches()
      outputs = dataset_utils.maybe_pad_batch(
          complete_batch,
          train=True,
          batch_size=desired_bs,
          pixel_level=pixel_level,
      )
      # Check output shape:
      self.assertEqual(outputs["inputs"].shape, (desired_bs, 28, 28, 3))
      # Check batch_mask:
      self.assertEqual(
          outputs["batch_mask"].shape, (desired_bs,) + batch_mask_shape
      )
      if pre_padding_mask:
          self.assertEqual(
              outputs["batch_mask"].sum(), complete_batch_mask.sum()
          )
      else:
          self.assertEqual(
              outputs["batch_mask"].sum(),
              float(desired_bs * mask_sum_coef),
          )
  
      print(complete_batch.keys())
      ######## Test partial training batches
      # Assert that the code throws an error as we dont handle partial training
      # batches in the codebase:
      with self.assertRaises(ValueError):
          _ = dataset_utils.maybe_pad_batch(
              partial_batch,
              train=True,
              batch_size=desired_bs,
              pixel_level=pixel_level,
          )
  
      ######## Test complete test batches
      (
          complete_batch,
          partial_batch,
          complete_batch_mask,
          partial_batch_mask,
      ) = make_fake_batches()
      outputs = dataset_utils.maybe_pad_batch(
          complete_batch,
          train=False,
          batch_size=desired_bs,
          pixel_level=pixel_level,
      )
      # Check output shape:
      self.assertEqual(outputs["inputs"].shape, (desired_bs, 28, 28, 3))
      # Check batch_mask:
      self.assertEqual(
          outputs["batch_mask"].shape, (desired_bs,) + batch_mask_shape
      )
      if pre_padding_mask:
          self.assertEqual(
              outputs["batch_mask"].sum(), complete_batch_mask.sum()
          )
      else:
          self.assertEqual(
              outputs["batch_mask"].sum(),
              float(desired_bs * mask_sum_coef),
          )
  
      ######## Test partial test batches
      outputs = dataset_utils.maybe_pad_batch(
          partial_batch,
          train=False,
          batch_size=desired_bs,
          pixel_level=pixel_level,
      )
      # Check output shape:
      self.assertEqual(outputs["inputs"].shape, (desired_bs, 28, 28, 3))
  
      # check output padding
      expected_out_pad = jnp.array(
          np.zeros((desired_bs - partial_bs, 28, 28, 3))
      )
      out_pad = outputs["inputs"][partial_bs:, :, :, :]
      self.assertTrue(jnp.array_equal(out_pad, expected_out_pad))
  
      # Check batch_mask:
      self.assertEqual(
          outputs["batch_mask"].shape, (desired_bs,) + batch_mask_shape
      )
  
      batch_mask = jnp.concatenate(
          [
              jnp.array(np.ones((partial_bs,) + batch_mask_shape)),
              jnp.array(
                  np.zeros(
                      (desired_bs - partial_bs,) + batch_mask_shape
                  )
              ),
          ],
          axis=0,
      )
      if pre_padding_mask:
          padded_pre_padding_mask = jnp.concatenate(
              [
                  partial_batch_mask,
                  jnp.array(
                      np.zeros(
                          (desired_bs - partial_bs,) + batch_mask_shape
                      )
                  ),
              ],
              axis=0,
          )
          batch_mask *= padded_pre_padding_mask
  
      self.assertTrue(
          jnp.array_equal(outputs["batch_mask"], batch_mask)
      )
  
file_path: scenic/scenic/dataset_lib/tests/test_dataset_utils.py
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  import jax.numpy
  import numpy
  from scenic.dataset_lib import dataset_utils
label_desc: Use image node_id 1 for calling the DatasetUtilsTest obj's underlying member method code with example usage: obj.test_maybe_pad_batch(pixel_level, batch_mask_shape, mask_sum_coef, pre_padding_mask) and returns: complete_batch,  partial_batch,  complete_batch_mask,  partial_batch_mask