class_docstr: Tests Dataset Utilities.
class_docstr_tok: ["Tests","Dataset","Utilities","."]
docstr: Tests mixup augmentation for different input formats and numpys.
docstr_tok: ["Tests","mixup","augmentation","for","different","input","formats","and","numpys","."]
code: def test_mixup(self, image_format, inputs_shape, jax_numpy):
      """Tests mixup augmentation for different input formats and numpys."""
      bs = inputs_shape[0]
      num_classes = 10
  
      if jax_numpy:
          np_backend = jnp
          mixup_fn = jax.jit(
              functools.partial(
                  dataset_utils.mixup,
                  alpha=1.0,
                  image_format=image_format,
                  rng=jax.random.PRNGKey(0),
              )
          )
      else:
          np_backend = np
          mixup_fn = functools.partial(
              dataset_utils.mixup,
              alpha=1.0,
              image_format=image_format,
              rng=None,
          )
  
      # Make a fake batch:
      inputs = np_backend.array(
          np.concatenate(
              (
                  np.zeros(shape=(bs // 2,) + inputs_shape[1:]),
                  np.ones(shape=(bs // 2,) + inputs_shape[1:]),
              ),
              axis=0,
          )
      )
      labels = np_backend.array(
          jax.nn.one_hot(
              np.concatenate(
                  (
                      np.ones(shape=(bs // 2,)),  # class 1
                      np.ones(shape=(bs // 2,)) * 2,  # class 2
                  ),
                  axis=0,
              ),
              num_classes,
          )
      )
      fake_batch = {"inputs": inputs, "label": labels}
  
      # Apply mixup:
      mixedup_batch = mixup_fn(fake_batch)
  
      self.assertEqual(mixedup_batch["inputs"].shape, inputs_shape)
      self.assertEqual(mixedup_batch["label"].shape, (bs, num_classes))
  
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 2 for calling the DatasetUtilsTest obj's underlying member method code with example usage: obj.test_mixup(image_format, inputs_shape, jax_numpy) without return types