code: def _get_inputs(
      key, same_inputs, shape, fn=jnp.cos
  ) -> tuple[jnp.ndarray, jnp.ndarray]:
      key, split = random.split(key)
      x1 = fn(random.normal(key, shape))
      batch_axis = shape.index(BATCH_SIZE)
      shape = (
          shape[:batch_axis]
          + (2 * BATCH_SIZE,)
          + shape[batch_axis + 1 :]
      )
      x2 = None if same_inputs else fn(random.normal(split, shape)) * 2
      return x1, x2
  
imports: import functools
  import itertools
  import random
  from absl.testing import absltest
  from jax import config
  from jax import default_backend
  from jax import jit
  from jax import random
  from jax.example_libraries import stax
  import jax.numpy
  import neural_tangents
  from neural_tangents import stax
  from neural_tangents._src.empirical import _DEFAULT_TESTING_NTK_IMPLEMENTATION
  from tests import test_utils
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: _get_inputs(key, same_inputs, shape, fn) and returns: x1,  x2