code: def test_strict_promotion(self, module_fn: ModuleFn, shape, dtype):
      f = hk.transform_with_state(
          lambda x: module_fn()(x)
      )  # pylint: disable=unnecessary-lambda
      rng = jax.random.PRNGKey(42)
      x = np.ones(shape, dtype)
      params, state = jax.eval_shape(f.init, rng, x)
      self.assertIsNotNone(
          jax.eval_shape(f.apply, params, state, rng, x)
      )
  
file_path: dm-haiku/haiku/_src/integration/rank_promotion_test.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import haiku
  from haiku._src import test_utils
  from haiku._src.integration import descriptors
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 3 for calling the RankPromotionTest obj's underlying member method code with example usage: obj.test_strict_promotion(module_fn, shape, dtype) without return types