code: def _label_data(
      rb_transitions: types.Transition,
      demonstration_transitions: types.Transition,
      mixup_alpha: Optional[float],
      key: networks_lib.PRNGKey,
  ) -> Tuple[types.Transition, jnp.ndarray]:
      """Create a tuple data, labels by concatenating the rb and dem transitions."""
      data = tree.map_structure(
          lambda x, y: jnp.concatenate([x, y]),
          rb_transitions,
          demonstration_transitions,
      )
      labels = jnp.concatenate(
          [
              jnp.zeros(rb_transitions.reward.shape),
              jnp.ones(demonstration_transitions.reward.shape),
          ]
      )
  
      if mixup_alpha is not None:
          lambda_key, mixup_key = jax.random.split(key)
  
          lambdas = tfd.Beta(mixup_alpha, mixup_alpha).sample(
              len(labels), seed=lambda_key
          )
  
          shuffled_data = tree.map_structure(
              lambda x: jax.random.permutation(key=mixup_key, x=x), data
          )
          shuffled_labels = jax.random.permutation(
              key=mixup_key, x=labels
          )
  
          data = tree.map_structure(
              lambda x, y: _weighted_average(x, y, lambdas),
              data,
              shuffled_data,
          )
          labels = _weighted_average(labels, shuffled_labels, lambdas)
  
      return data, labels
  
file_path: acme/acme/agents/jax/ail/losses.py
imports: import functools
  from typing import Callable, Dict, Optional, Tuple
  from acme import types
  from acme.jax import networks
  import jax
  import jax.numpy
  import tensorflow_probability
  import tree
label_desc: Use image node_id 3 for calling a global function with example usage: _label_data(rb_transitions, demonstration_transitions, mixup_alpha, key) and returns: data,  labels