code: def _num_layers_and_params(params):
      """Returns number of weight layers and total number of params in params."""
      leaves = jax.tree_util.tree_leaves(params)
      return len(leaves), sum(
          [jnp.prod(jnp.array(v.shape)) for v in leaves]
      )
  
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 a global function with example usage: _num_layers_and_params(params) and returns: len,  sum