code: def _get_dummy_data(
      mask_constant: float,
  ) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]:
      """Return dummy data for when downloading embeddings is not feasible."""
      n_train, n_test = 6, 6
  
      def get_x(shape, key):
          key_x, key_mask = random.split(key)
          x = random.normal(key_x, shape)
          mask = random.bernoulli(key_mask, 0.6, shape)
          x = jnp.where(mask, mask_constant, x)
          return x
  
      def get_y(x):
          x = jnp.where(x == mask_constant, 0.0, x)
  
          def weighted_sum(x, start, end):
              return jnp.sum(
                  x[..., start:end]
                  * jnp.arange(x.shape[1])[None, ..., None],
                  axis=(1, 2),
              )
  
          y_label = (
              jnp.stack(
                  [
                      weighted_sum(x, 0, x.shape[-1] // 2),
                      weighted_sum(x, x.shape[-1] // 2, x.shape[-1]),
                  ],
                  axis=-1,
              )
              > 0
          )
          y = jnp.where(y_label, 0.5, -0.5)
          return y
  
      rng_train, rng_test = random.split(random.PRNGKey(1), 2)
      x_train = get_x((n_train, _MAX_SENTENCE_LENGTH, 50), rng_train)
      x_test = get_x((n_test, _MAX_SENTENCE_LENGTH, 50), rng_test)
  
      y_train, y_test = get_y(x_train), get_y(x_test)
      return x_train, y_train, x_test, y_test
  
file_path: neural-tangents/examples/imdb.py
imports: import time
  from absl import app
  from jax import random
  import jax.numpy
  import neural_tangents
  from neural_tangents import stax
  from examples import datasets
  from examples import util
label_desc: Use image node_id 2 for calling a global function with example usage: _get_dummy_data(mask_constant) and returns: x_train,  y_train,  x_test,  y_test, x, y, jnp