code: def test_sonnet_module_as_layer(self, module_fn, input_shape, dtype):
      mod = module_fn()
      layer = LayerAdapter(module=module_fn(), dtype=dtype)
      example_input = tf.ones(input_shape, dtype=dtype)
  
      # Check outputs are the same.
      for m_y, l_y in zip(
          tree.flatten(mod(example_input)),
          tree.flatten(layer(example_input)),
      ):
          self.assertEqual(m_y.shape, l_y.shape)
          self.assertEqual(m_y.dtype, l_y.dtype)
  
      # Some modules (e.g. Sequential) are parameter-less.
      snt.allow_empty_variables(mod)
  
      # Check that variables are the same.
      self.assertEqual(len(mod.variables), len(layer.variables))
      self.assertEqual(
          len(mod.trainable_variables), len(layer.trainable_variables)
      )
  
      # Check that Keras layer freezing works
      layer.trainable = False
      self.assertEmpty(layer.trainable_variables)
  
imports: from absl.testing import parameterized
  import sonnet
  from sonnet.src import test_utils
  from sonnet.src.conformance import descriptors
  import tensorflow
  import tree
label_desc: Use image node_id 2 for calling the KerasTest obj's underlying member method code with example usage: obj.test_sonnet_module_as_layer(module_fn, input_shape, dtype) without return types