class_docstr: Tests for ViT+ layers.
code: def test_mlp_block(self, num_layers, expected_num_weights):
      batch_size, in_dim, out_dim = 8, 20, 30
      mlp_configs = {"out_dim": out_dim, "num_layers": num_layers}
      mlp = layers.PredictorMLP(**mlp_configs)
      inputs = jnp.ones((batch_size, in_dim))
      out, variables = mlp.init_with_output(
          jax.random.PRNGKey(0), inputs
      )
      self.assertEqual(out.shape, (batch_size, out_dim))
      weights, _ = jax.tree_util.tree_flatten(variables)
      self.assertLen(weights, expected_num_weights)
  
file_path: scenic/scenic/projects/owl_vit/tests/layers_test.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  import jax.numpy
  import ml_collections
  from scenic.projects.owl_vit import layers
label_desc: Use image node_id 1 for calling the LayersTest obj's underlying member method code with example usage: obj.test_mlp_block(num_layers, expected_num_weights) without return types