code: def test_mlp(self, dropout_rate, output_size, num_layers):
      # Input definition.
      nodes = self.graphs.nodes
  
      # Model definition.
      mlp = models.MLP(
          feature_sizes=[output_size] * num_layers,
          dropout_rate=dropout_rate,
          activation=lambda x: x,
          deterministic=False,
      )
      nodes_after_mlp, _ = mlp.init_with_output(self.rngs, nodes)
  
      # Test that dropout actually worked.
      num_masked_entries = jnp.sum(nodes_after_mlp == 0)
      num_total_entries = jnp.size(nodes_after_mlp)
      self.assertLessEqual(
          num_masked_entries, (dropout_rate + 0.05) * num_total_entries
      )
      self.assertLessEqual(
          (dropout_rate - 0.05) * num_total_entries, num_masked_entries
      )
  
      # Test the shape of the output.
      self.assertEqual(nodes_after_mlp.shape[-1], output_size)
  
file_path: flax/examples/ogbg_molpcba/models_test.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  import jax.numpy
  import jraph
  import models
label_desc: Use image node_id 2 for calling the ModelsTest obj's underlying member method code with example usage: obj.test_mlp(dropout_rate, output_size, num_layers) without return types