code: def test_primitives_by_source(self):
      def f(x, y):
          s = jnp.sin(x)
          return jnp.sin(s) + jnp.cos(y)
  
      hist = jaxpr_util.primitives_by_source(
          make_jaxpr(f)(1.0, 1.0).jaxpr
      )
  
      sin_keys = [k for k in hist.keys() if k.startswith("sin @ ")]
      rem_keys = [k for k in hist.keys() if not k.startswith("sin @ ")]
  
      self.assertEqual(sum(hist[k] for k in sin_keys), 2)
      self.assertTrue(all(hist[k] == 1 for k in rem_keys))
  
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 2 for calling the JaxprStatsTest obj's underlying member method code with example usage: obj.test_primitives_by_source() and returns: unknown