code: def get_structure(
      eqn: Optional[JaxprEqn],
      invals: list[Union[ShapedArray, AbstractValue]],
      idx: int,
      _s_rules: bool,
  ) -> Structure:
      if any(i is AbstractValue for i in invals):
          raise TypeError(invals)
  
      if eqn is None:
          # Identity function
          primitive = None
          cts_in = invals[0]
          assert idx == 0
  
      else:
          if len(eqn.outvars) != 1:
              raise NotImplementedError(eqn)
          cts_in = eqn.outvars[0].aval
  
          primitive = eqn.primitive
          assert len(invals) == len(eqn.invars)
          assert 0 <= idx < len(eqn.invars)
  
      if not isinstance(cts_in, ShapedArray):
          raise TypeError(cts_in)
  
      if primitive in STRUCTURE_RULES and _s_rules:
          structure = STRUCTURE_RULES[primitive](
              eqn, idx, invals, cts_in
          )
  
      else:
          # No simplification rule found.
          structure = Structure()
  
      # TODO(romann): can we avoid special-casing `reshape`s?
      if primitive == lax.reshape_p:
          cts_in = ShapedArray(invals[idx].shape, invals[idx].dtype)
  
      # Check that number of trace output and input axes match.
      assert len(structure.in_trace) == len(structure.out_trace)
  
      # Check that input and output traced sizes are the same.
      out_trace_size = utils.size_at(cts_in, structure.out_trace)
      in_trace_size = utils.size_at(invals[idx], structure.in_trace)
      assert in_trace_size == out_trace_size
  
      # Check that number of input/output diagonal axes match.
      assert len(structure.out_diagonal) == len(structure.in_diagonal)
  
      # Check for each output diagonal axis there's only input axes of correct
      # size or `None`. Inval axis should be not `None`.
      for out_d, in_d in zip(
          structure.out_diagonal, structure.in_diagonal
      ):
          assert len(in_d) == len(invals)
          assert in_d[idx] is not None
          for ix, i in enumerate(in_d):
              if i is not None:
                  assert invals[ix].shape[i] == cts_in.shape[out_d]
  
      return 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 1 for calling a global function with example usage: get_structure(eqn, invals, idx, _s_rules) and returns: structure