code: def test_primitives_by_shape(self):
      def f(x, y):
          def sub(x, y):
              return jnp.sum(jnp.array([x, y])), y
  
          s, _ = jit(sub)(x, y)
          return jnp.sin(s) + jnp.cos(y)
  
      hist = jaxpr_util.primitives_by_shape(
          make_jaxpr(f)(1.0, 1.0).jaxpr
      )
  
      t = "64" if config.enable_x64.value else "32"
      shapes = [
          f"add :: float{t}[]",
          f"sin :: float{t}[]",
          f"cos :: float{t}[]",
          f"reduce_sum :: float{t}[]",
          f"concatenate :: float{t}[2]",
          f"pjit :: float{t}[] *",
      ]
      for k in shapes:
          self.assertEqual(hist[k], 1)
  
file_path: jax/tests/jaxpr_util_test.py
imports: import os
  import gzip
  import json
  from absl.testing import absltest
  import jax
  from jax import jit, make_jaxpr, numpy
  from jax._src import config
  from jax._src import jaxpr_util
  from jax._src import test_util
  from jax._src.lib import xla_client
label_desc: Use image node_id 3 for calling the JaxprStatsTest obj's underlying member method code with example usage: obj.test_primitives_by_shape() and returns: unknown, jnp,  y