code: def _CheckShapeAgreement(test_case, init_fun, apply_fun, input_shape):
      rng_key = random.PRNGKey(0)
      rng_key, init_key = random.split(rng_key)
      result_shape, params = init_fun(init_key, input_shape)
      inputs = random_inputs(test_case.rng(), input_shape)
      if params:
          inputs = inputs.astype(np.dtype(params[0]))
      result = apply_fun(params, inputs, rng=rng_key)
      test_case.assertEqual(result.shape, result_shape)
  
imports: from absl.testing import absltest
  import numpy
  from jax._src import test_util
  from jax import random
  from jax.example_libraries import stax
  from jax import dtypes
  from jax import config
label_desc: Use image node_id 2 for calling a global function with example usage: _CheckShapeAgreement(test_case, init_fun, apply_fun, input_shape) without return types