code: def test_primitives(self):
      def f(x, y):
          s = jit(jnp.sin)(x)
          return jnp.sin(s) + jnp.cos(y)
  
      hist = jaxpr_util.primitives(make_jaxpr(f)(1.0, 1.0).jaxpr)
  
      primitives = ["add", "sin", "cos"]
      primitives.append("pjit")
      for k in primitives:
          assert k in hist, k
      self.assertEqual(hist["sin"], 2)
      self.assertTrue(
          all(count == 1 for k, count in hist.items() if k != "sin")
      )
  
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 1 for calling the JaxprStatsTest obj's underlying member method code with example usage: obj.test_primitives() and returns: unknown