code: def primitives_by_shape(jaxpr: core.Jaxpr):
      def shape_fmt(var):
          return (
              "*"
              if isinstance(var, core.DropVar)
              else var.aval.str_short()
          )
  
      def key(eqn):
          return (
              eqn.primitive.name,
              " ".join(map(shape_fmt, eqn.outvars)),
          )
  
      return histogram(jaxpr, key, " :: ".join)
  
imports: from __future__ import annotations
  from collections import Counter, defaultdict
  import gzip
  import itertools
  import json
  import types
  from typing import Any, Callable, Union
  from jax._src import core
  from jax._src import util
  from jax._src import source_info_util
  from jax._src.lib import xla_client
label_desc: Use image node_id 6 for calling a global function with example usage: primitives_by_shape(jaxpr) and returns: histogram, unknown, eqn,  str