class_docstr: A `Feature` that tracks shape information in a graph.
  
  This `Feature` aids in the replacement of all `Shape`\s and `Subtensor`\s of `Shape`\s with
  `Shape_i` and `MakeVector` `Op`\s.
  
  This `Feature` and its associated rewrites have several goals:
  
  1. to "lift" `Shape`\s to as close to the inputs as possible,
  2. to infer the shape of every node in the graph in terms of the
     input shapes, and
  3. remove fill `Op`\s (e.g. `Second`) from the graph.
  
  Lifting shapes as close to the inputs as possible is important for
  canonicalization because it is very bad form to have to compute
  something just to know how big it will be.  Firstly, it is a waste
  of time to compute such outputs.  But it is important to get rid
  of these outputs as early as possible in the compilation process
  because the extra computations make it appear as if many internal
  graph nodes have multiple clients.  Many rewrites refuse to
  work on nodes with multiple clients.
  
  Lifting is done by using an `<Op>.infer_shape` function if one is
  present, or else using a conservative default.  An Op that
  supports shape-lifting should define a infer_shape(self, fgraph, node,
  input_shapes) function.  The argument input_shapes is a tuple of
  tuples... there is an interior tuple for each input to the node.
  The tuple has as many elements as dimensions.  The element in
  position i of tuple j represents the i'th shape component of the
  j'th input.  The function should return a tuple of tuples.  One
  output tuple for each node.output.  Again, the i'th element of the
  j'th output tuple represents the output[j].shape[i] of the
  function.  If an output is not a TensorType, then None should be
  returned instead of a tuple for that output.
  
  For example the infer_shape for a matrix-matrix product would accept
  input_shapes=((x0,x1), (y0,y1)) and return ((x0, y1),).
  
  Inferring the shape of internal nodes in the graph is important
  for doing size-driven rewrites.  If we know how big various
  intermediate results will be, we can estimate the cost of many Ops
  accurately, and generate c-code that is specific [e.g. unrolled]
  to particular sizes.
  
  In cases where you cannot figure out the shape, raise a ShapeError.
  
  Notes
  -----
  To use this shape information in rewrites, use the
  ``shape_of`` dictionary.
  
  For example:
  
  .. code-block:: python
  
      try:
          shape_of = fgraph.shape_feature.shape_of
      except AttributeError:
          # This can happen when the mode doesn't include the ShapeFeature.
          return
  
      shape_of_output_zero = shape_of[node.output[0]]
  
  The ``shape_of_output_zero`` symbol will contain a tuple, whose
  elements are either integers or symbolic integers.
  
  TODO: check to see if the symbols are necessarily
  non-constant... or are integer literals sometimes PyTensor
  constants?? That would be confusing.
class_docstr_tok: ["A","`","Feature","`","that","tracks","shape","information","in","a","graph",".","This","`","Feature","`","aids","in","the","replacement","of","all","`","Shape","`","\\s","and","`","Subtensor","`","\\s","of","`","Shape","`","\\s","with","`","Shape_i","`","and","`","MakeVector","`","`","Op","`","\\s",".","This","`","Feature","`","and","its","associated","rewrites","have","several","goals",":","1.","to","``","lift","''","`","Shape","`","\\s","to","as","close","to","the","inputs","as","possible",",","2.","to","infer","the","shape","of","every","node","in","the","graph","in","terms","of","the","input","shapes",",","and","3.","remove","fill","`","Op","`","\\s","(","e.g",".","`","Second","`",")","from","the","graph",".","Lifting","shapes","as","close","to","the","inputs","as","possible","is","important","for","canonicalization","because","it","is","very","bad","form","to","have","to","compute","something","just","to","know","how","big","it","will","be",".","Firstly",",","it","is","a","waste","of","time","to","compute","such","outputs",".","But","it","is","important","to","get","rid","of","these","outputs","as","early","as","possible","in","the","compilation","process","because","the","extra","computations","make","it","appear","as","if","many","internal","graph","nodes","have","multiple","clients",".","Many","rewrites","refuse","to","work","on","nodes","with","multiple","clients",".","Lifting","is","done","by","using","an","`","<","Op",">",".infer_shape","`","function","if","one","is","present",",","or","else","using","a","conservative","default",".","An","Op","that","supports","shape-lifting","should","define","a","infer_shape","(","self",",","fgraph",",","node",",","input_shapes",")","function",".","The","argument","input_shapes","is","a","tuple","of","tuples","...","there","is","an","interior","tuple","for","each","input","to","the","node",".","The","tuple","has","as","many","elements","as","dimensions",".","The","element","in","position","i","of","tuple","j","represents","the","i'th","shape","component","of","the","j'th","input",".","The","function","should","return","a","tuple","of","tuples",".","One","output","tuple","for","each","node.output",".","Again",",","the","i'th","element","of","the","j'th","output","tuple","represents","the","output","[","j","]",".shape","[","i","]","of","the","function",".","If","an","output","is","not","a","TensorType",",","then","None","should","be","returned","instead","of","a","tuple","for","that","output",".","For","example","the","infer_shape","for","a","matrix-matrix","product","would","accept","input_shapes=","(","(","x0",",","x1",")",",","(","y0",",","y1",")",")","and","return","(","(","x0",",","y1",")",",",")",".","Inferring","the","shape","of","internal","nodes","in","the","graph","is","important","for","doing","size-driven","rewrites",".","If","we","know","how","big","various","intermediate","results","will","be",",","we","can","estimate","the","cost","of","many","Ops","accurately",",","and","generate","c-code","that","is","specific","[","e.g",".","unrolled","]","to","particular","sizes",".","In","cases","where","you","can","not","figure","out","the","shape",",","raise","a","ShapeError",".","Notes","--","--","-","To","use","this","shape","information","in","rewrites",",","use","the","``","shape_of","``","dictionary",".","For","example",":","..","code-block",":",":","python","try",":","shape_of","=","fgraph.shape_feature.shape_of","except","AttributeError",":","#","This","can","happen","when","the","mode","does","n't","include","the","ShapeFeature",".","return","shape_of_output_zero","=","shape_of","[","node.output","[","0","]","]","The","``","shape_of_output_zero","``","symbol","will","contain","a","tuple",",","whose","elements","are","either","integers","or","symbolic","integers",".","TODO",":","check","to","see","if","the","symbols","are","necessarily","non-constant","...","or","are","integer","literals","sometimes","PyTensor","constants","?","?","That","would","be","confusing","."]
code: def on_import(self, fgraph, node, reason):
      if node.outputs[0] in self.shape_of:
          # this is a revert, not really an import
          for r in node.outputs + node.inputs:
              assert r in self.shape_of
          return
  
      for i, r in enumerate(node.inputs):
          # make sure we have shapes for the inputs
          self.init_r(r)
  
      o_shapes = self.get_node_infer_shape(node)
  
      # this is packed information
      # an element of o_shapes is either None or a tuple
      #   elements of the tuple can be either strings, or ints
      if len(o_shapes) != len(node.outputs):
          raise Exception(
              f'The infer_shape method for the Op "{node.op}" returned a list '
              f"with the wrong number of element: len(o_shapes) = {len(o_shapes)} "
              f" != len(node.outputs) = {len(node.outputs)}"
          )
  
      # Ensure shapes are in 'int64'. This is to make sure the assert
      # found in the `local_useless_subtensor` rewrite does not fail.
      for sh_idx, sh in enumerate(o_shapes):
          if sh is None:
              continue
          if not isinstance(sh, (list, tuple)):
              raise ValueError(
                  f"infer_shape of {node} didn't return a list of"
                  f" list. It returned '{o_shapes}'"
              )
          new_shape = []
          for i, d in enumerate(sh):
              # Note: we ignore any shape element that is not typed (i.e.,
              # does not have a 'dtype' attribute). This means there may
              # still remain int elements that are int32 on 32-bit platforms,
              # but this works with `local_useless_subtensor`, so for now we
              # keep it this way. See #266 for a better long-term fix.
              if getattr(d, "dtype", "int64") != "int64":
                  assert d.dtype in discrete_dtypes, (node, d.dtype)
                  assert str(d.dtype) != "uint64", node
                  new_shape += sh[len(new_shape) : i + 1]
                  if isinstance(d, Constant):
                      casted_d = constant(d.data, dtype="int64")
                  else:
                      casted_d = cast(d, "int64")
                  new_shape[i] = casted_d
          if new_shape:
              # We replace the shape with wrong dtype by the one with
              # 'int64'.
              new_shape += sh[len(new_shape) :]
              o_shapes[sh_idx] = tuple(new_shape)
  
      for r, s in zip(node.outputs, o_shapes):
          self.set_shape(r, s)
  
imports: import traceback
  from io import StringIO
  from typing import Optional
  from typing import cast
  from warnings import warn
  import numpy
  import pytensor
  from pytensor.configdefaults import config
  from pytensor.graph.basic import Constant, Variable, ancestors, equal_computations
  from pytensor.graph.features import AlreadyThere, Feature
  from pytensor.graph.fg import FunctionGraph
  from pytensor.graph.rewriting.basic import GraphRewriter, check_chain, copy_stack_trace, node_rewriter
  from pytensor.graph.utils import InconsistencyError, get_variable_trace_string
  from pytensor.tensor.basic import MakeVector, as_tensor_variable, cast, constant, extract_constant, get_underlying_scalar_constant_value, register_infer_shape, stack
  from pytensor.tensor.elemwise import DimShuffle, Elemwise
  from pytensor.tensor.exceptions import NotScalarConstantError, ShapeError
  from pytensor.tensor.rewriting.basic import register_canonicalize, register_specialize, register_stabilize, register_useless, topo_constant_folding
  from pytensor.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape, Unbroadcast, shape_i, specify_shape, unbroadcast
  from pytensor.tensor.subtensor import Subtensor, get_idx_list
  from pytensor.tensor.type import TensorType, discrete_dtypes, integer_dtypes
  from pytensor.tensor.type_other import NoneConst
label_desc: Use image node_id 14 for calling the ShapeFeature obj's underlying member method code with example usage: obj.on_import(fgraph, node, reason) without return types