code: def get_structure_cache(
      jaxpr: Jaxpr, _s_rules: bool
  ) -> dict[Var, Structure]:
      """Associates a least common structure to each input variable of the `jaxpr`.
  
      Args:
        jaxpr: Jaxpr to build cache for.
        _s_rules: whether to use structure rules or not.
  
      Returns:
        A dictionary mapping input variables to the least common structure of all
        primitives it is present in as a direct input.
      """
      invar_to_structure: dict[Var, Structure] = {}
  
      for var in jaxpr.invars:
          if var in jaxpr.outvars:
              if isinstance(var, Literal):
                  raise TypeError(var)
  
              # Identity function
              structure = get_id_structure(var.aval, _s_rules)
  
              if var in invar_to_structure:
                  invar_to_structure[var] &= structure
              else:
                  invar_to_structure[var] = structure
  
      for eqn in jaxpr.eqns:
          for i_eqn, var in enumerate(eqn.invars):
              if var in jaxpr.invars:
                  if isinstance(var, Literal):
                      raise TypeError(var)
  
                  structure = get_structure(
                      eqn=eqn,
                      invals=[v.aval for v in eqn.invars],
                      idx=i_eqn,
                      _s_rules=_s_rules,
                  )
  
                  if var in invar_to_structure:
                      invar_to_structure[var] &= structure
                  else:
                      invar_to_structure[var] = structure
  
      return invar_to_structure
  
imports: import functools
  from typing import Any, Callable, Optional, Union
  import jax
  from jax import lax
  from jax.core import AbstractValue
  from jax.core import Jaxpr
  from jax.core import JaxprEqn
  from jax.core import Literal
  from jax.core import Primitive
  from jax.core import ShapedArray
  from jax.core import Var
  from jax.interpreters import ad
  import jax.numpy
  import numpy
  from .None import utils
  from .dataclasses import dataclass
  from .dataclasses import field
label_desc: Use image node_id 2 for calling a global function with example usage: get_structure_cache(jaxpr, _s_rules) and returns: invar_to_structure