code: def generate_parity_eval_sample(batch_size, seq_len):
      """Generate one sample for parity task.
  
      Args:
        batch_size: Determines the batch size.
        seq_len: Determines the sequence length of parity vector.
  
      Yields:
        One Sample.
  
      """
      rng = jax.random.PRNGKey(42)
      while True:
          rng, _ = jax.random.split(rng)
          # Parity:
          sample = jax.random.choice(
              rng,
              a=jnp.array((1.0, 0.0, -1.0), jnp.float32),
              shape=(batch_size, seq_len),
          )
          label = (
              jnp.sum(jnp.equal(sample, 1.0), axis=-1).astype(jnp.int32)
              % 2
          )
          sample = jax.nn.one_hot(sample + 1, 3).astype(jnp.float32)
          yield {"inputs": sample, "label": label}
  
imports: import functools
  import jax
  import jax.numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib import datasets
label_desc: Use image node_id 2 for calling a global function with example usage: generate_parity_eval_sample(batch_size, seq_len) without return types