code: def _CheckShapeAgreement(test_case, init_fun, apply_fun, input_shape):
      jax_rng = random.PRNGKey(0)
      result_shape, params = init_fun(jax_rng, input_shape)
      result = apply_fun(
          params,
          test_case.rng.normal(size=input_shape).astype("float32"),
      )
      test_case.assertEqual(result.shape, result_shape)
  
file_path: jax/examples/examples_test.py
imports: import os
  import sys
  import zlib
  from absl.testing import absltest
  from absl.testing import parameterized
  import numpy
  from jax import lax
  from jax import random
  import jax.numpy
  from examples import kernel_lsq
  from jax import config
label_desc: Use image node_id 1 for calling a global function with example usage: _CheckShapeAgreement(test_case, init_fun, apply_fun, input_shape) without return types