code: """This module provides the `Scan` `Op`.
  
  Memory reuse in scan
  --------------------
  
  To reduce the number of memory allocations and copies associated with calling
  the inner function and recovering the outputs at every iteration, Scan uses a
  memory pre-allocation mechanism for some of its outputs. Instead of repeatedly
  calling the inner function and copying the outputs to designated locations,
  it tries to make the inner function write the outputs directly to the
  designated locations.
  
  This is achieved by initializing, at every iteration, the output storage
  of the inner function with references to previously allocated memory. Other
  than the code in the Python and Cython backends to do this and to ensure that
  the pre-allocated memory has been used, the memory pre-allocation mechanism
  relies on the following elements to work properly :
  - In make_thunk(), when compiling the inner function, the borrow flag must
    be set to False for the inputs. This will prevent aliasing between the
    inputs and the outputs of the inner function which could lead to invalid
    results.
  - In make_thunk(), again, the borrow flag must be set to True for the outputs.
    This will make Aesara consider the output storages as persistent and make
    Aesara provide them as pre-allocated storage to the ops that compute the
    outputs of the inner function instead of letting these ops allocate their
    own output storage.
  - The ops that produce the outputs of the inner function must be prevented
    from working inplace because if they do, they're not using the pre-allocated
    storage. This is achieved by including the optimization
    'add_no_output_from_inplace' to the compilation mode used by scan. It
    prevents other optimizations from altering the graph such that outputs are
    produced by inplace operations.
  - The ScanSaveMem optimization, whose goal is to limit the amount of memory
    used by scan, needs to allocate buffers large enough to be able, at every
    iteration, to simultaneously read the needed previous states and storing
    the new states. Before the memory reuse feature, the buffers could be
    smaller because, often, Scan only needed buffers large enough to read the
    needed previous states. This is because all the outputs of the inner
    function were computed before any of them was stored in the buffers. Now,
    the outputs are stored as they are computed which means that, if the buffer
    is too small, computing an output can overwrite an input that is still
    needed to compute another output.
  
  """
  
  
  import copy
  import itertools
  import logging
  import time
  from collections import OrderedDict
  
  import numpy as np
  
  import aesara
  from aesara import tensor as aet
  from aesara.compile.builders import infer_shape
  from aesara.compile.function import function
  from aesara.compile.io import In, Out
  from aesara.compile.mode import AddFeatureOptimizer, get_mode
  from aesara.compile.profiling import ScanProfileStats, register_profiler_printer
  from aesara.configdefaults import config
  from aesara.gradient import DisconnectedType, NullType, Rop, grad, grad_undefined
  from aesara.graph.basic import (
      Apply,
      Constant,
      Variable,
      clone_replace,
      equal_computations,
      graph_inputs,
      io_connection_pattern,
  )
  from aesara.graph.features import NoOutputFromInplace
  from aesara.graph.fg import MissingInputError
  from aesara.graph.op import Op, ops_with_inner_function
  from aesara.link.c.basic import CLinker
  from aesara.link.c.exceptions import MissingGXX
  from aesara.link.utils import raise_with_op
  from aesara.scan.utils import Validator, forced_replace, hash_listsDictsTuples, safe_new
  from aesara.tensor.basic import as_tensor_variable
  from aesara.tensor.math import minimum
  from aesara.tensor.shape import Shape_i
  from aesara.tensor.type import TensorType, integer_dtypes
  from aesara.tensor.var import TensorVariable
  
  
  __docformat__ = "restructedtext en"
  __authors__ = (
      "Razvan Pascanu "
      "Frederic Bastien "
      "James Bergstra "
      "Pascal Lamblin "
      "PyMC Developers "
      "Aesara Developers "
  )
  __copyright__ = "(c) 2010, Universite de Montreal"
  
  # Logging function for sending warning or info
  _logger = logging.getLogger("aesara.scan.op")
  
  
  class Scan(Op):
      """
  
      Parameters
      ----------
      inputs
          Inputs of the inner function of scan.
      outputs
          Outputs of the inner function of scan.
      info
          Dictionary containing different properties of the scan op (like number
          of different types of arguments, name, mode, if it should run on GPU or
          not, etc.).
      typeConstructor
          Function that constructs an equivalent to Aesara TensorType.
  
      Notes
      -----
      ``typeConstructor`` had been added to refactor how
      Aesara deals with the GPU. If it runs on the GPU, scan needs
      to construct certain outputs (those who reside in the GPU
      memory) as the GPU-specific type.  However we can not import
      gpu code in this file (as it is in sandbox, and not available
      on each machine) so the workaround is that the GPU
      optimization passes to the constructor of this class a
      function that is able to construct a GPU type. This way the
      class Scan does not need to be aware of the details for the
      GPU, it just constructs any tensor using this function (which
      by default constructs normal tensors).
  
      """
  
      def __init__(
          self,
          inputs,
          outputs,
          info,
          typeConstructor=None,
      ):
          # adding properties into self
          self.inputs = inputs
          self.outputs = outputs
          self.__dict__.update(info)
          # I keep a version of info in self, to use in __eq__ and __hash__,
          # since info contains all tunable parameters of the op, so for two
          # scan to be equal this tunable parameters should be the same
          self.info = info
          # build a list of output types for any Apply node using this op.
          self.output_types = []
          idx = 0
          jdx = 0
  
          def tensorConstructor(broadcastable, dtype):
              return TensorType(broadcastable=broadcastable, dtype=dtype)
  
          if typeConstructor is None:
              typeConstructor = tensorConstructor
  
          while idx < self.n_mit_mot_outs:
              # Not that for mit_mot there are several output slices per
              # output sequence
              o = outputs[idx]
              self.output_types.append(
                  typeConstructor(
                      broadcastable=(False,) + o.type.broadcastable, dtype=o.type.dtype
                  )
              )
  
              idx += len(self.mit_mot_out_slices[jdx])
              jdx += 1
  
          # mit_sot / sit_sot / nit_sot
          end = idx + self.n_mit_sot + self.n_sit_sot + self.n_nit_sot
  
          for o in outputs[idx:end]:
              self.output_types.append(
                  typeConstructor(
                      broadcastable=(False,) + o.type.broadcastable, dtype=o.type.dtype
                  )
              )
  
          # shared outputs + possibly the ending condition
          for o in outputs[end:]:
              self.output_types.append(o.type)
  
          if self.as_while:
              self.output_types = self.output_types[:-1]
  
          mode_instance = get_mode(self.mode)
          # Clone mode_instance, altering "allow_gc" for the linker,
          # and adding a message if we profile
          if self.name:
              message = self.name + " sub profile"
          else:
              message = "Scan sub profile"
  
          self.mode_instance = mode_instance.clone(
              link_kwargs=dict(allow_gc=self.allow_gc), message=message
          )
  
          if not hasattr(self, "name") or self.name is None:
              self.name = "scan_fn"
          # to have a fair __eq__ comparison later on, we update the info with
          # the actual mode used to compile the function and the name of the
          # function that we set in case none was given
          self.info["name"] = self.name
  
          # Pre-computing some values to speed up perform
          self.mintaps = [np.min(x) for x in self.tap_array]
          self.mintaps += [0 for x in range(self.n_nit_sot)]
          self.seqs_arg_offset = 1 + self.n_seqs
          self.shared_arg_offset = (
              self.seqs_arg_offset + self.n_mit_mot + self.n_mit_sot + self.n_sit_sot
          )
          self.nit_sot_arg_offset = self.shared_arg_offset + self.n_shared_outs
          self.n_outs = self.n_mit_mot + self.n_mit_sot + self.n_sit_sot
          self.n_tap_outs = self.n_mit_mot + self.n_mit_sot
          if self.info["gpua"]:
              self._hash_inner_graph = self.info["gpu_hash"]
          else:
              # Do the missing inputs check here to have the error early.
              for var in graph_inputs(self.outputs, self.inputs):
                  if var not in self.inputs and not isinstance(var, Constant):
                      raise MissingInputError(f"ScanOp is missing an input: {repr(var)}")
              self._cmodule_key = CLinker().cmodule_key_variables(
                  self.inputs, self.outputs, []
              )
              self._hash_inner_graph = hash(self._cmodule_key)
  
          # Compute mappings between outer inputs, outer outputs, inner
          # inputs and inner outputs to determine with variables are associated
          # with the same states.
          self.var_mappings = self.get_oinp_iinp_iout_oout_mappings()
  
      def validate_inner_graph(self):
          """
          Perform some elementary validations on the inner graph to ensure
          that it is coherent.
  
          """
  
          # For every recurrent output, iterate over the associated inner
          # inputs and output and ensure that they have the same dtype
          nb_recurr_outputs = self.n_mit_mot + self.n_mit_sot + self.n_sit_sot
  
          for outer_oidx in range(nb_recurr_outputs):
  
              inner_iidxs = self.var_mappings["inner_inp_from_outer_out"][outer_oidx]
              inner_oidxs = self.var_mappings["inner_out_from_outer_out"][outer_oidx]
  
              for (inner_iidx, inner_oidx) in itertools.product(inner_iidxs, inner_oidxs):
  
                  type_input = self.inputs[inner_iidx].type
                  type_output = self.outputs[inner_oidx].type
                  if type_input != type_output:
                      raise TypeError(
                          "Inconsistency in the inner graph of "
                          f"scan '{self.name}' : an input and an output are "
                          "associated with the same recurrent state "
                          "and should have the same type but have "
                          f"type '{type_input}' and '{type_output}' respectively."
                      )
  
          # If scan has the flag 'gpua' set to false (meaning that is shouldn't
          # use the gpuarray gpu backend ), ensure that is has no input and no
          # output with type GpuArrayType
          from aesara.gpuarray import GpuArrayType
  
          if not self.info.get("gpua", False):
              for inp in self.inputs:
                  if isinstance(inp.type, GpuArrayType):
                      raise TypeError(
                          "Inconsistency in the inner graph of "
                          f"scan '{self.name}' : one of the inputs to the "
                          "inner graph is of type GpuArrayType but "
                          "the attributes of the scan op indicate "
                          "that it shouldn't be the case"
                      )
  
              for out in self.outputs:
                  if isinstance(out.type, GpuArrayType):
                      raise TypeError(
                          "Inconsistency in the inner graph of "
                          f"scan '{self.name}' : one of the outputs to the "
                          "inner graph is of type GpuArrayType but "
                          "the attributes of the scan op indicate "
                          "that it shouldn't be the case"
                      )
  
      def __setstate__(self, d):
          self.__dict__.update(d)
          if "allow_gc" not in self.__dict__:
              self.allow_gc = True
              self.info["allow_gc"] = True
          if not hasattr(self, "var_mappings"):
              # Generate the mappings between inner and outer inputs and outputs
              # if they haven't already been generated.
              self.var_mappings = self.get_oinp_iinp_iout_oout_mappings()
          if hasattr(self, "fn"):
              if not hasattr(self, "thunk_mit_mot_out_slices"):
                  # The thunk has been compiled before mit_mot preallocation
                  # feature was implemented. Mark every mit_mot output tap as
                  # not having been preallocated
                  self.mitmots_preallocated = [False] * self.n_mit_mot_outs
  
              if not hasattr(self, "outs_is_tensor"):
                  # The thunk has been compiled before the analysis, at
                  # compilation time, of the location of the inputs and outputs.
                  # Perform this analysis here.
                  self.inps_is_tensor = [
                      isinstance(out, TensorVariable)
                      for out in self.fn.maker.fgraph.inputs
                  ]
                  self.outs_is_tensor = [
                      isinstance(out, TensorVariable)
                      for out in self.fn.maker.fgraph.outputs
                  ]
  
          # Ensure that the graph associated with the inner function is valid.
          self.validate_inner_graph()
  
      def make_node(self, *inputs):
          """
          Conventions:
              inner_X - the variable corresponding to X in the inner function
                        of scan (the lambda function executed at every time
                        step)
              outer_X - the variable corresponding to X in the outer graph,
                        i.e. the main graph (where the scan op lives)
              inner_X_out - the variable representing the new value of X after
                            executing one step of scan (i.e. outputs given by
                            the inner function)
  
          """
          assert np.all(isinstance(i, Variable) for i in inputs)
          # Check that the number of inputs to the Scan node corresponds to
          # the number of inputs of the inner function of scan
          n_outer_ins = len(inputs) - len(self.outer_nitsot(inputs)) - 1
          n_inner_ins = (
              len(self.inner_seqs(self.inputs))
              + len(self.mitmot_taps())
              + len(self.mitsot_taps())
              + len(self.inner_sitsot(self.inputs))
              + len(self.inner_shared(self.inputs))
              + len(self.inner_non_seqs(self.inputs))
          )
          assert n_outer_ins == n_inner_ins, (
              "The number of inputs given to the inner function of scan"
              " does not match the number of inputs given to scan."
          )
          # Force the inputs to be on the CPU
          new_inputs = [as_tensor_variable(inputs[0])]
          # assert dtype is consistent
          err_msg1 = (
              "When compiling the inner function of scan (the "
              "function called by scan in each of its iterations) "
              "the following error has been encountered: The "
              "%s %s (argument number %d) has dtype "
              "%s and %d dimension(s). The corresponding variable "
              "in the inner function of scan %s "
              "however has dtype %s and %d dimension(s). This "
              "variable in the inner function of scan should "
              "have the same dtype and one fewer dimension "
              "compared to its corresponding variable in the initial "
              "state (outputs_info in scan nomenclature). For example, "
              "if the inner function of scan returns a vector "
              "of size d and scan uses the values of "
              "the previous time-step, then the initial state in scan "
              "should be a matrix of shape (1, d). "
              "The first dimension of this "
              "matrix corresponds to the number of previous time-steps "
              "that scan uses in each of its iterations. "
              "In order to solve this issue if the two variable currently "
              "have the same dimensionality, you can increase the "
              "dimensionality of the varialbe in the initial state of scan "
              "by using dimshuffle or shape_padleft. "
          )
          err_msg2 = (
              "When compiling the inner function of scan the "
              "following error has been encountered: The "
              "initial state (`outputs_info` in scan nomenclature) "
              "of variable %s (argument number %d) "
              "has dtype %s, while the result of the inner function "
              "(`fn`) has dtype %s. This can happen if the inner "
              "function of scan results in an upcast or downcast."
          )
          err_msg3 = (
              "When compiling the inner function of scan (the "
              "function called by scan in each of its iterations) "
              "the following error has been encountered: The "
              "initial state (`outputs_info` in scan nomenclature) "
              "of variable %s (argument number %d) has %d dimension(s), "
              "while the corresponding variable in the result of the inner "
              "function of scan (`fn`) has %d dimension(s) (it should "
              "be one less than the initial state). For example, "
              "if the inner function of scan returns a vector "
              "of size d and scan uses the values of "
              "the previous time-step, then the initial state in scan "
              "should be a matrix of shape (1, d). "
              "The first dimension of this "
              "matrix corresponds to the number of previous time-steps "
              "that scan uses in each of its iterations. "
              "In order to solve this issue if the two varialbe currently "
              "have the same dimensionality, you can increase the "
              "dimensionality of the variable in the initial state of scan "
              "by using dimshuffle or shape_padleft. "
          )
  
          def check_broadcast(v1, v2):
              """Checks that the broadcast pattern of v1 and v2.
  
              Controls that the broadcast pattern of the variable provided as
              input to `scan` matches the broadcast pattern provided in
              `output_info`. It raises an error when they don't match. The
              typical case is when the user provides either the input or the
              `output_info` (but not both) with a dimension fixed to 1,
              which may wrongly be interpreted as broadcastable.
  
              """
              if not hasattr(v1, "broadcastable") and not hasattr(v2, "broadcastable"):
                  return
              msg = (
                  "The broadcast pattern of the output of scan (%s) is "
                  "inconsistent with the one provided in `output_info` "
                  "(%s). The output on axis %d is `%r`, but it is `%r` on "
                  "axis %d in `output_info`. This can happen if one of the "
                  "dimension is fixed to 1 in the input, while it is still "
                  "variable in the output, or vice-verca. You have to make "
                  "them consistent, e.g. using aesara.tensor."
                  "{patternbroadcast,unbroadcast,addbroadcast}."
              )
              size = min(len(v1.broadcastable), len(v2.broadcastable))
              for n, (b1, b2) in enumerate(
                  zip(v1.broadcastable[-size:], v2.broadcastable[-size:])
              ):
                  if b1 != b2:
                      a1 = n + size - len(v1.broadcastable) + 1
                      a2 = n + size - len(v2.broadcastable) + 1
                      raise TypeError(msg % (v1.type, v2.type, a1, b1, b2, a2))
  
          def format(var, as_var):
              """
              This functions ensures that ``out`` has the same dtype as
              ``inp`` as well as calling filter_variable to make sure
              they are both TensorType or GpuArrayType. It internally
              deals with the corner case where inp.ndim + 1 = out.ndim
  
              """
              if not hasattr(var, "dtype"):
                  return var
              rval = var
              if rval.type.dtype != as_var.type.dtype:
                  rval = rval.astype(as_var.type.dtype)
              if rval.ndim == as_var.ndim:
                  rval = as_var.type.filter_variable(rval)
              else:
                  tmp = as_var.type.clone(
                      broadcastable=(
                          tuple(var.broadcastable[:1]) + tuple(as_var.broadcastable)
                      )
                  )
                  rval = tmp.filter_variable(rval)
              return rval
  
          # Check if input sequences and variables representing a slice of
          # them have the same dtype
          argoffset = 0
          for inner_seq, outer_seq in zip(
              self.inner_seqs(self.inputs), self.outer_seqs(inputs)
          ):
              check_broadcast(outer_seq, inner_seq)
              new_inputs.append(format(outer_seq, as_var=inner_seq))
  
          argoffset += len(self.outer_seqs(inputs))
          # Check that this 3 things have the same dtype for mit_mot:
          #   - initial state of the output
          #   - variable representing an input slice of the output
          #   - variable representing an output slice of the output
          ipos = 0
          opos = 0
          inner_mitmot = self.inner_mitmot(self.inputs)
          inner_mitmot_outs = self.inner_mitmot_outs(self.outputs)
          for idx, (itaps, otaps, _outer_mitmot) in enumerate(
              zip(self.mitmot_taps(), self.mitmot_out_taps(), self.outer_mitmot(inputs))
          ):
              outer_mitmot = format(_outer_mitmot, as_var=inner_mitmot[ipos])
              new_inputs.append(outer_mitmot)
              for k in range(len(itaps)):
                  if (
                      inner_mitmot[ipos + k].type.dtype != outer_mitmot.type.dtype
                      or inner_mitmot[ipos + k].ndim != outer_mitmot.ndim - 1
                  ):
                      raise ValueError(
                          err_msg1
                          % (
                              "initial state (outputs_info" " in scan nomenclature) ",
                              str(outer_mitmot),
                              argoffset + idx,
                              outer_mitmot.type.dtype,
                              outer_mitmot.type.ndim,
                              str(inner_mitmot[ipos + k]),
                              inner_mitmot[ipos + k].type.dtype,
                              inner_mitmot[ipos + k].type.ndim,
                          )
                      )
              ipos += len(itaps)
              for k in range(len(otaps)):
                  if inner_mitmot_outs[opos + k].type.dtype != outer_mitmot.type.dtype:
                      raise ValueError(
                          err_msg2
                          % (
                              str(outer_mitmot),
                              argoffset + idx,
                              outer_mitmot.type.dtype,
                              inner_mitmot_outs[opos + k].type.dtype,
                          )
                      )
                  if inner_mitmot_outs[opos + k].ndim != outer_mitmot.ndim - 1:
                      raise ValueError(
                          err_msg3
                          % (
                              str(outer_mitmot),
                              argoffset + idx,
                              outer_mitmot.ndim,
                              inner_mitmot_outs[opos + k].ndim,
                          )
                      )
              opos += len(otaps)
          argoffset += len(self.outer_mitmot(inputs))
          # Same checks as above but for outputs of type mit_sot
          ipos = 0
          inner_mitsots = self.inner_mitsot(self.inputs)
          for idx, (itaps, _outer_mitsot, inner_mitsot_out) in enumerate(
              zip(
                  self.mitsot_taps(),
                  self.outer_mitsot(inputs),
                  self.inner_mitsot_outs(self.outputs),
              )
          ):
              outer_mitsot = format(_outer_mitsot, as_var=inner_mitsots[ipos])
              new_inputs.append(outer_mitsot)
  
              for k in range(len(itaps)):
                  if (
                      inner_mitsots[ipos + k].type.dtype != outer_mitsot.type.dtype
                      or inner_mitsots[ipos + k].ndim != outer_mitsot.ndim - 1
                  ):
                      raise ValueError(
                          err_msg1
                          % (
                              "initial state (outputs_info" " in scan nomenclature) ",
                              str(outer_mitsot),
                              argoffset + idx,
                              outer_mitsot.type.dtype,
                              outer_mitsot.type.ndim,
                              str(inner_mitsots[ipos + k]),
                              inner_mitsots[ipos + k].type.dtype,
                              inner_mitsots[ipos + k].type.ndim,
                          )
                      )
              ipos += len(itaps)
              if inner_mitsot_out.type.dtype != outer_mitsot.type.dtype:
                  raise ValueError(
                      err_msg2
                      % (
                          str(outer_mitsot),
                          argoffset + idx,
                          outer_mitsot.type.dtype,
                          inner_mitsot_out.type.dtype,
                      )
                  )
              if inner_mitsot_out.ndim != outer_mitsot.ndim - 1:
                  raise ValueError(
                      err_msg3
                      % (
                          str(outer_mitsot),
                          argoffset + idx,
                          outer_mitsot.ndim,
                          inner_mitsot_out.ndim,
                      )
                  )
  
          argoffset += len(self.outer_mitsot(inputs))
          # Same checks as above but for outputs of type sit_sot
          for idx, (inner_sitsot, _outer_sitsot, inner_sitsot_out) in enumerate(
              zip(
                  self.inner_sitsot(self.inputs),
                  self.outer_sitsot(inputs),
                  self.inner_sitsot_outs(self.outputs),
              )
          ):
              outer_sitsot = format(_outer_sitsot, as_var=inner_sitsot)
              new_inputs.append(outer_sitsot)
              if inner_sitsot.ndim != outer_sitsot.ndim - 1:
                  raise ValueError(
                      err_msg1
                      % (
                          "initial state (outputs_info" " in scan nomenclature) ",
                          str(outer_sitsot),
                          argoffset + idx,
                          outer_sitsot.type.dtype,
                          outer_sitsot.type.ndim,
                          str(inner_sitsot),
                          inner_sitsot.type.dtype,
                          inner_sitsot.type.ndim,
                      )
                  )
              if inner_sitsot_out.type.dtype != outer_sitsot.type.dtype:
                  raise ValueError(
                      err_msg2
                      % (
                          str(outer_sitsot),
                          argoffset + idx,
                          outer_sitsot.type.dtype,
                          inner_sitsot_out.type.dtype,
                      )
                  )
              if inner_sitsot_out.ndim != outer_sitsot.ndim - 1:
                  raise ValueError(
                      err_msg3
                      % (
                          str(outer_sitsot),
                          argoffset + idx,
                          outer_sitsot.type.ndim,
                          inner_sitsot_out.type.ndim,
                      )
                  )
  
          argoffset += len(self.outer_sitsot(inputs))
          # Check that the shared variable and their update rule have the same
          # dtype. Maybe even same type ?!
          for idx, (inner_shared, inner_shared_out, _outer_shared) in enumerate(
              zip(
                  self.inner_shared(self.inputs),
                  self.inner_shared_outs(self.outputs),
                  self.outer_shared(inputs),
              )
          ):
              outer_shared = format(_outer_shared, as_var=inner_shared)
              new_inputs.append(outer_shared)
              if (
                  hasattr(outer_shared, "dtype")
                  and outer_shared.dtype != inner_shared_out.dtype
              ):
                  raise ValueError(
                      err_msg2
                      % (
                          str(outer_shared),
                          idx + argoffset,
                          outer_shared.dtype,
                          inner_shared_out.dtype,
                      )
                  )
              if (
                  hasattr(outer_shared, "dtype")
                  and outer_shared.ndim != inner_shared_out.ndim
              ):
                  raise ValueError(
                      err_msg3
                      % (
                          str(outer_shared),
                          idx + argoffset,
                          outer_shared.ndim,
                          inner_shared_out.ndim,
                      )
                  )
  
              if hasattr(outer_shared, "dtype") and (
                  outer_shared.dtype != inner_shared.dtype
                  or outer_shared.ndim != inner_shared.ndim
              ):
                  raise ValueError(
                      err_msg1
                      % (
                          "initial state (outputs_info" " in scan nomenclature) ",
                          str(outer_shared),
                          argoffset + idx,
                          outer_shared.dtype,
                          outer_shared.ndim,
                          str(inner_shared),
                          inner_shared.dtype,
                          inner_shared.ndim,
                      )
                  )
          # We do not need to call `format` on outer_nisot arguments.
          # outer_nitsot stands for no input tap single output tap. This means
          # these are states that do not feed anything back in the recurrent
          # computation, and hence they do not have an initial state. The scan
          # node however receives an input for each such argument, the input
          # in this case is just a int saying how many steps of this output we
          # need to store. This input does not have the same dtype, nor is it the same
          # type of tensor as the output, it is always a scalar int.
          new_inputs += [as_tensor_variable(ons) for ons in self.outer_nitsot(inputs)]
          for inner_nonseq, _outer_nonseq in zip(
              self.inner_non_seqs(self.inputs), self.outer_non_seqs(inputs)
          ):
              outer_nonseq = format(_outer_nonseq, as_var=inner_nonseq)
              new_inputs.append(outer_nonseq)
              if inner_nonseq.type != outer_nonseq.type:
                  raise ValueError(
                      (
                          "Argument %s given to scan node does not"
                          " match its correspondence %s"
                      )
                      % (str(outer_nonseq), str(inner_nonseq))
                  )
  
          for outer_nitsot in self.outer_nitsot(inputs):
              # For every nit_sot input we get as input a int/uint that
              # depicts the size in memory for that sequence. This feature is
              # used by truncated BPTT and by scan space optimization
              if (
                  str(outer_nitsot.type.dtype) not in integer_dtypes
                  or outer_nitsot.ndim != 0
              ):
                  raise ValueError(
                      "For output %s you need to provide a " "scalar int !",
                      str(outer_nitsot),
                  )
          assert len(new_inputs) == len(inputs)
  
          # The vector_seqs and vector_outs are just a workaround
          # strange NumPy behavior: vector_ndarray[int] return a NumPy
          # scalar and not a NumPy ndarray of 0 dimensions.
          def is_cpu_vector(s):
              return isinstance(s.type, TensorType) and s.ndim == 1
  
          self.vector_seqs = [
              is_cpu_vector(seq) for seq in new_inputs[1 : 1 + self.n_seqs]
          ]
          self.vector_outs = [
              is_cpu_vector(arg)
              for arg in new_inputs[1 + self.n_seqs : (1 + self.n_seqs + self.n_outs)]
          ]
          self.vector_outs += [
              isinstance(t.type, TensorType) and t.ndim == 0
              for t in self.outer_nitsot_outs(self.outputs)
          ]
  
          apply_node = Apply(self, new_inputs, [t() for t in self.output_types])
          return apply_node
  
      def __eq__(self, other):
          # Check if we are dealing with same type of objects
          if not type(self) == type(other):
              return False
          if "destroy_map" not in self.info:
              self.info["destroy_map"] = OrderedDict()
          if "destroy_map" not in other.info:
              other.info["destroy_map"] = OrderedDict()
          keys_to_check = [
              "truncate_gradient",
              "profile",
              "n_seqs",
              "tap_array",
              "as_while",
              "n_mit_sot",
              "destroy_map",
              "n_nit_sot",
              "n_shared_outs",
              "n_sit_sot",
              "gpua",
              "n_mit_mot_outs",
              "n_mit_mot",
              "mit_mot_out_slices",
          ]
          # This are some safety checks ( namely that the inner graph has the
          # same number of inputs and same number of outputs )
          if not len(self.inputs) == len(other.inputs):
              return False
          elif not len(self.outputs) == len(other.outputs):
              return False
          for key in keys_to_check:
              if self.info[key] != other.info[key]:
                  return False
          # If everything went OK up to here, there is still one thing to
          # check. Namely, do the internal graph represent same
          # computations
          for self_in, other_in in zip(self.inputs, other.inputs):
              if self_in.type != other_in.type:
                  return False
  
          return equal_computations(
              self.outputs, other.outputs, self.inputs, other.inputs
          )
  
      def __str__(self):
          if self.gpua:
              gpu_str = "gpu"
          else:
              gpu_str = "cpu"
          if self.as_while:
              name = "do_while"
          else:
              name = "for"
          aux_txt = "%s"
          if len(self.destroy_map.keys()) > 0:
              # Check if all outputs are inplace
              if sorted(self.destroy_map.keys()) == sorted(
                  range(self.n_mit_mot + self.n_mit_sot + self.n_sit_sot)
              ):
                  aux_txt += "all_inplace,%s,%s}"
              else:
                  aux_txt += "{inplace{"
                  for k in self.destroy_map.keys():
                      aux_txt += str(k) + ","
                  aux_txt += "},%s,%s}"
          else:
              aux_txt += "{%s,%s}"
          aux_txt = aux_txt % (name, gpu_str, str(self.name))
          return aux_txt
  
      def __hash__(self):
          return hash(
              (
                  type(self),
                  # and a hash representing the inner graph using the
                  # CLinker.cmodule_key_
                  self._hash_inner_graph,
                  hash_listsDictsTuples(self.info),
              )
          )
  
      def make_thunk(self, node, storage_map, compute_map, no_recycling, impl=None):
          """
  
          Parameters
          ----------
          node
              Something previously returned by self.make_node.
          storage_map
              dict variable -> one-element-list where a computed
              value for this variable may be found.
          compute_map
              dict variable -> one-element-list where a boolean
              value will be found. The boolean indicates whether the
              variable's storage_map container contains a valid value (True)
              or if it has not been computed yet (False).
          no_recycling
              List of variables for which it is forbidden to reuse memory
              allocated by a previous call.
          impl
              Use 'py' if we want python execution.
          Notes
          -----
          If the thunk consults the storage_map on every call, it is safe
          for it to ignore the no_recycling argument, because elements of the
          no_recycling list will have a value of None in the storage map. If
          the thunk can potentially cache return values (like CLinker does),
          then it must not do so for variables in the no_recycling list.
  
          """
  
          # Before building the thunk, validate that the inner graph is
          # coherent
          self.validate_inner_graph()
  
          # Setting up all my variables in what I believe is a more Cython
          # friendly form
  
          node_input_storage = [storage_map[r] for r in node.inputs]
          node_output_storage = [storage_map[r] for r in node.outputs]
          # If a shared variable is the result of a ViewOp it is a clear
          # indication that we need to copy that value after the perform of
          # scan is done
          slices = self.n_mit_mot_outs + self.n_mit_sot + self.n_sit_sot + self.n_nit_sot
  
          if config.scan__allow_output_prealloc:
  
              # Go through the mitmots. Whenever a mitmot has a tap both as an
              # input and an output, wrap the input such that the corresponding
              # output variable becomes an update to be performed on it, possibly
              # inplace at the end of the functions's execution.
              wrapped_inputs = [In(x, borrow=False) for x in self.inputs[: self.n_seqs]]
              new_outputs = [x for x in self.outputs]
              preallocated_mitmot_outs = []
              new_mit_mot_out_slices = copy.deepcopy(self.mit_mot_out_slices)
  
              input_idx = self.n_seqs
              for mitmot_idx in range(self.n_mit_mot):
                  for inp_tap in self.tap_array[mitmot_idx]:
                      if inp_tap in self.mit_mot_out_slices[mitmot_idx]:
                          inp = self.inputs[input_idx]
  
                          # Figure out the index of the corresponding output
                          output_idx = sum(
                              [len(m) for m in self.mit_mot_out_slices[:mitmot_idx]]
                          )
                          output_idx += self.mit_mot_out_slices[mitmot_idx].index(inp_tap)
  
                          # Make it so the input is automatically updated to the
                          # output value, possibly inplace, at the end of the
                          # function execution. Also, since an update is
                          # defined, a default value must also be (this is
                          # verified by DebugMode). Use an array of size 0 but
                          # the right ndim and dtype (use a shape of 1 on
                          # broadcastable dimensions, 0 on the others).
                          default_shape = [1 if _b else 0 for _b in inp.broadcastable]
                          default_val = inp.type.value_zeros(default_shape)
                          wrapped_inp = In(
                              variable=inp,
                              value=default_val,
                              update=self.outputs[output_idx],
                          )
                          wrapped_inputs.append(wrapped_inp)
                          preallocated_mitmot_outs.append(output_idx)
                          new_mit_mot_out_slices[mitmot_idx].remove(inp_tap)
                      else:
                          # Wrap the corresponding input as usual. Leave the
                          # output as-is.
                          wrapped_inputs.append(In(self.inputs[input_idx], borrow=False))
                      input_idx += 1
  
              # Wrap the inputs not associated to mitmots and wrap the remaining
              # outputs
              wrapped_inputs += [In(x, borrow=False) for x in self.inputs[input_idx:]]
              wrapped_outputs = [Out(x, borrow=True) for x in new_outputs[:slices]]
              wrapped_outputs += new_outputs[slices:]
  
              # Remove now useless outputs from the output list (start from the
              # end to avoid altering the indices of the other outputs to be
              # deleted.
              preallocated_mitmot_outs.sort()
              for p in preallocated_mitmot_outs[::-1]:
                  del wrapped_outputs[p]
  
              # Store the list of mitmot output taps that have been altered
              # so they can be preallocated
              self.mitmots_preallocated = [
                  i in preallocated_mitmot_outs for i in range(self.n_mit_mot_outs)
              ]
  
              # Add an optimization to the compilation mode to attach a feature
              # to the function graph just before the inplace optimizations are
              # applied (inplace optimizations start at position 50 so the
              # optimization to attach the feature is registered at position 49.9
              # so that it runs before them). This feature will prevent mitsot,
              # sitsot and nitsot outputs from being computed inplace (to allow
              # their preallocation).
              mitsot_start = self.n_mit_mot_outs - len(preallocated_mitmot_outs)
              nitsot_end = mitsot_start + self.n_mit_sot + self.n_sit_sot + self.n_nit_sot
              feature = NoOutputFromInplace(mitsot_start, nitsot_end)
              opt = AddFeatureOptimizer(feature)
              compilation_mode = self.mode_instance.register((opt, 49.9))
  
          else:
              # Output preallocation is not activated. Mark every mitmot output
              # tap as not being preallocated
              self.mitmots_preallocated = [False] * self.n_mit_mot_outs
  
              wrapped_inputs = [In(x, borrow=True) for x in self.inputs]
              wrapped_outputs = [Out(x, borrow=False) for x in self.outputs[:slices]]
              wrapped_outputs += self.outputs[slices:]
  
              compilation_mode = self.mode_instance
  
          profile = None
          if config.profile or (
              isinstance(self.profile, (str, bool, (int,))) and self.profile
          ):
              if isinstance(self.profile, str):
                  profile = ScanProfileStats(name=self.profile)
              else:
                  profile = ScanProfileStats(name=self.name)
          elif self.profile:
              profile = self.profile
          # make_thunk can be called many times on the same op
          # we do not want to recompile the inner fct every time.
          if not getattr(self, "fn", None):
              self.fn = function(
                  wrapped_inputs,
                  wrapped_outputs,
                  mode=compilation_mode,
                  name=self.name,
                  profile=profile,
                  on_unused_input="ignore",
              )
  
          # Analyse the compile inner function to determine which inputs and
          # outputs are on the gpu and speed up some checks during the execution
          self.inps_is_tensor = [
              isinstance(out, TensorVariable) for out in self.fn.maker.fgraph.inputs
          ]
          self.outs_is_tensor = [
              isinstance(out, TensorVariable) for out in self.fn.maker.fgraph.outputs
          ]
  
          try:
              if impl == "py":
                  raise MissingGXX
              cython_mintaps = np.asarray(self.mintaps, dtype="int32")
              cython_tap_array_len = np.asarray(
                  [len(x) for x in self.tap_array], dtype="int32"
              )
              if len(self.tap_array) == 0:
                  d1 = 0
              else:
                  d1 = np.max(cython_tap_array_len)
              d0 = len(self.tap_array)
              cython_tap_array = np.zeros((d0, d1), dtype="int32")
              for _d0 in range(d0):
                  for _d1 in range(cython_tap_array_len[_d0]):
                      cython_tap_array[_d0, _d1] = self.tap_array[_d0][_d1]
              cython_mit_mot_out_nslices = np.asarray(
                  [len(x) for x in self.mit_mot_out_slices], dtype="int32"
              )
              if len(self.mit_mot_out_slices) == 0:
                  d1 = 0
              else:
                  d1 = np.max(cython_mit_mot_out_nslices)
              d0 = len(self.mit_mot_out_slices)
              cython_mit_mot_out_slices = np.zeros((d0, d1), dtype="int32")
              for _d0 in range(d0):
                  for _d1 in range(cython_mit_mot_out_nslices[_d0]):
                      cython_mit_mot_out_slices[_d0, _d1] = self.mit_mot_out_slices[_d0][
                          _d1
                      ]
  
              cython_vector_seqs = np.asarray(self.vector_seqs, dtype="int32")
              cython_vector_outs = np.asarray(self.vector_outs, dtype="int32")
              cython_mitmots_preallocated = np.asarray(
                  self.mitmots_preallocated, dtype="int32"
              )
  
              cython_inps_is_tensor = np.asarray(self.inps_is_tensor, dtype="int32")
              cython_outs_is_tensor = np.asarray(self.outs_is_tensor, dtype="int32")
  
              if self.destroy_map:
                  cython_destroy_map = [
                      x in self.destroy_map for x in range(len(node.outputs))
                  ]
              else:
                  cython_destroy_map = [0 for x in range(len(node.outputs))]
              cython_destroy_map = np.asarray(cython_destroy_map, dtype="int32")
              from . import scan_perform_ext
  
              def p(node, args, outs):
                  return scan_perform_ext.perform(
                      self.n_shared_outs,
                      self.n_mit_mot_outs,
                      self.n_seqs,
                      self.n_mit_mot,
                      self.n_mit_sot,
                      self.n_sit_sot,
                      self.n_nit_sot,
                      args[0],
                      self.as_while,
                      cython_mintaps,
                      cython_tap_array,
                      cython_tap_array_len,
                      cython_vector_seqs,
                      cython_vector_outs,
                      cython_mit_mot_out_slices,
                      cython_mit_mot_out_nslices,
                      cython_mitmots_preallocated,
                      cython_inps_is_tensor,
                      cython_outs_is_tensor,
                      self.fn.fn,
                      self.fn,
                      cython_destroy_map,
                      args,
                      outs,
                      self,
                      node,
                  )
  
          except (ImportError, MissingGXX):
              p = self.perform
  
          # default arguments are stored in the closure of `rval`
  
          # Big ugly hack since we can't get the real value of allow_gc
          # for the englobing function.
          allow_gc = config.allow_gc and not self.allow_gc
  
          def rval(
              p=p, i=node_input_storage, o=node_output_storage, n=node, allow_gc=allow_gc
          ):
              r = p(n, [x[0] for x in i], o)
              for o in node.outputs:
                  compute_map[o][0] = True
              if allow_gc:
                  self.fn.free()
              return r
  
          rval.inputs = node_input_storage
          rval.outputs = node_output_storage
          rval.perform = p
          rval.lazy = False
          return rval
  
      def inner_seqs(self, list_inputs):
          # Given the list of inner inputs this function grabs those
          # corresponding to sequences
          return list_inputs[: self.n_seqs]
  
      def outer_seqs(self, list_inputs):
          if isinstance(list_inputs, Apply):
              list_inputs = list_inputs.inputs
          # Given the list of outer inputs this function grabs those
          # corresponding to sequences
          return list_inputs[1 : 1 + self.n_seqs]
  
      def inner_mitmot(self, list_inputs):
          n_taps = sum(len(x) for x in self.tap_array[: self.n_mit_mot])
          return list_inputs[self.n_seqs : self.n_seqs + n_taps]
  
      def outer_mitmot(self, list_inputs):
          if isinstance(list_inputs, Apply):
              list_inputs = list_inputs.inputs
          return list_inputs[1 + self.n_seqs : 1 + self.n_seqs + self.n_mit_mot]
  
      def inner_mitmot_outs(self, list_outputs):
          n_taps = sum(len(x) for x in self.mit_mot_out_slices)
          return list_outputs[:n_taps]
  
      def outer_mitmot_outs(self, list_outputs):
          if isinstance(list_outputs, Apply):
              list_outputs = list_outputs.outputs
          return list_outputs[: self.n_mit_mot]
  
      def mitmot_taps(self):
          return self.tap_array[: self.n_mit_mot]
  
      def mitmot_out_taps(self):
          return self.mit_mot_out_slices[: self.n_mit_mot]
  
      def inner_mitsot(self, list_inputs):
          n_mitmot_taps = sum(len(x) for x in self.tap_array[: self.n_mit_mot])
          ntaps_upto_sit_sot = sum(
              len(x) for x in self.tap_array[: (self.n_mit_mot + self.n_mit_sot)]
          )
          return list_inputs[
              self.n_seqs + n_mitmot_taps : self.n_seqs + ntaps_upto_sit_sot
          ]
  
      def outer_mitsot(self, list_inputs):
          if isinstance(list_inputs, Apply):
              list_inputs = list_inputs.inputs
          offset = 1 + self.n_seqs + self.n_mit_mot
          return list_inputs[offset : offset + self.n_mit_sot]
  
      def inner_mitsot_outs(self, list_outputs):
          n_taps = sum(len(x) for x in self.mit_mot_out_slices)
          return list_outputs[n_taps : n_taps + self.n_mit_sot]
  
      def outer_mitsot_outs(self, list_outputs):
          if isinstance(list_outputs, Apply):
              list_outputs = list_outputs.outputs
          return list_outputs[self.n_mit_mot : self.n_mit_mot + self.n_mit_sot]
  
      def mitsot_taps(self):
          return self.tap_array[self.n_mit_mot : self.n_mit_mot + self.n_mit_sot]
  
      def inner_sitsot(self, list_inputs):
          n_taps_upto_sit_sot = sum(
              len(x) for x in self.tap_array[: (self.n_mit_mot + self.n_mit_sot)]
          )
          offset = self.n_seqs + n_taps_upto_sit_sot
          return list_inputs[offset : offset + self.n_sit_sot]
  
      def outer_sitsot(self, list_inputs):
          if isinstance(list_inputs, Apply):
              list_inputs = list_inputs.inputs
          offset = 1 + self.n_seqs + self.n_mit_mot + self.n_mit_sot
          return list_inputs[offset : offset + self.n_sit_sot]
  
      def inner_sitsot_outs(self, list_outputs):
          n_taps = sum(len(x) for x in self.mit_mot_out_slices)
          offset = self.n_mit_sot + n_taps
          return list_outputs[offset : offset + self.n_sit_sot]
  
      def outer_sitsot_outs(self, list_outputs):
          if isinstance(list_outputs, Apply):
              list_outputs = list_outputs.outputs
          offset = self.n_mit_mot + self.n_mit_sot
          return list_outputs[offset : offset + self.n_sit_sot]
  
      def outer_nitsot(self, list_inputs):
          if isinstance(list_inputs, Apply):
              list_inputs = list_inputs.inputs
          offset = (
              1
              + self.n_seqs
              + self.n_mit_mot
              + self.n_mit_sot
              + self.n_sit_sot
              + self.n_shared_outs
          )
          return list_inputs[offset : offset + self.n_nit_sot]
  
      def inner_nitsot_outs(self, list_outputs):
          n_taps = sum(len(x) for x in self.mit_mot_out_slices)
          offset = self.n_mit_sot + n_taps + self.n_sit_sot
          return list_outputs[offset : offset + self.n_nit_sot]
  
      def outer_nitsot_outs(self, list_outputs):
          if isinstance(list_outputs, Apply):
              list_outputs = list_outputs.outputs
          offset = self.n_mit_mot + self.n_mit_sot + self.n_sit_sot
          return list_outputs[offset : offset + self.n_nit_sot]
  
      def inner_shared(self, list_inputs):
          n_taps_upto_sit_sot = sum(
              len(x) for x in self.tap_array[: (self.n_mit_mot + self.n_mit_sot)]
          )
          offset = self.n_seqs + n_taps_upto_sit_sot + self.n_sit_sot
          return list_inputs[offset : offset + self.n_shared_outs]
  
      def outer_shared(self, list_inputs):
          if isinstance(list_inputs, Apply):
              list_inputs = list_inputs.inputs
          offset = 1 + self.n_seqs + self.n_mit_mot + self.n_mit_sot + self.n_sit_sot
          return list_inputs[offset : offset + self.n_shared_outs]
  
      def inner_shared_outs(self, list_outputs):
          n_taps = sum(len(x) for x in self.mit_mot_out_slices)
          offset = self.n_mit_sot + n_taps + self.n_sit_sot + self.n_nit_sot
          return list_outputs[offset : offset + self.n_shared_outs]
  
      def outer_shared_outs(self, list_outputs):
          if isinstance(list_outputs, Apply):
              list_outputs = list_outputs.outputs
          offset = self.n_mit_mot + self.n_mit_sot + self.n_sit_sot + self.n_nit_sot
          return list_outputs[offset : offset + self.n_shared_outs]
  
      def inner_non_seqs(self, list_inputs):
          n_taps_upto_sit_sot = sum(
              len(x) for x in self.tap_array[: (self.n_mit_mot + self.n_mit_sot)]
          )
          offset = self.n_seqs + n_taps_upto_sit_sot + self.n_sit_sot + self.n_shared_outs
          return list_inputs[offset:]
  
      def outer_non_seqs(self, list_inputs):
          if isinstance(list_inputs, Apply):
              list_inputs = list_inputs.inputs
          offset = (
              1
              + self.n_seqs
              + self.n_mit_mot
              + self.n_mit_sot
              + self.n_sit_sot
              + self.n_nit_sot
              + self.n_shared_outs
          )
          return list_inputs[offset:]
  
      def perform(self, node, inputs, output_storage, params=None):
          """Compute the scan operation in Python.
  
          The `inputs` are packed like this:
  
              n_steps
  
              X sequence inputs x_1, x_2, ... x_<self.n_seqs>
  
              Y initial states (u_1, u_2, ... u_<self.n_outs>) for our
              outputs. Each must have appropriate length (T_1, T_2, ..., T_Y).
  
              W other inputs w_1, w_2, ... w_W
  
          There are at least ``1 + self.n_seqs + self.n_outs`` inputs, and the
          ones above this number are passed to the scanned function as
          non-sequential inputs.
  
          The outputs are more straightforward:
  
              Y sequence outputs y_1, y_2, ... y_<self.n_outs>
  
          """
          # 1. Unzip the number of steps and sequences. If number of steps is
          # negative flip sequences around, and make n_steps positive
          t0_call = time.time()
          t_fn = 0
          n_steps = inputs[0]
          seqs = []
          if n_steps < 0:
              # History, in the past, this was used for backward
              # scan. Now we reverse the inputs outside of scan.
              raise IndexError(
                  f"Scan was asked to run for negative number of step {int(n_steps)}"
              )
          elif n_steps == 0:
              raise NotImplementedError(
                  "We didn't implemented yet the case where scan do 0 iteration"
              )
          else:
              for idx, seq in enumerate(inputs[1 : self.seqs_arg_offset]):
                  if seq.shape[0] < n_steps:
                      raise ValueError(
                          (
                              "Sequence is shorter then the required "
                              "number of steps : (n_steps, seq, "
                              "seq.shape):"
                          ),
                          n_steps,
                          node.inputs[1 + idx],
                          seq.shape,
                      )
                  seqs.append(seq)
  
          # 2. Allocate memory for the outputs. Construct the list:
          #       store_steps  -- map containing the length of each output
          #       pos          -- map containing the current position of each
          #                       output
  
          store_steps = [
              arg.shape[0]
              for arg in inputs[self.seqs_arg_offset : self.shared_arg_offset]
          ]
          store_steps += [
              arg
              for arg in inputs[
                  self.nit_sot_arg_offset : self.nit_sot_arg_offset + self.n_nit_sot
              ]
          ]
  
          pos = [
              (-self.mintaps[idx]) % store_steps[idx]
              for idx in range(self.n_outs + self.n_nit_sot)
          ]
          # 2.1 Create storage space for outputs
          for idx in range(self.n_outs):
              if idx in self.destroy_map:
                  # ^ Case 1. Outputs should be computed inplace of their
                  # initial state
                  output_storage[idx][0] = inputs[self.seqs_arg_offset + idx]
              elif (
                  output_storage[idx][0] is not None
                  and output_storage[idx][0].shape[1:]
                  == inputs[self.seqs_arg_offset + idx].shape[1:]
                  and output_storage[idx][0].shape[0] >= store_steps[idx]
              ):
                  # Put in the values of the initial state
                  output_storage[idx][0] = output_storage[idx][0][: store_steps[idx]]
                  if idx > self.n_mit_mot:
                      l = -self.mintaps[idx]
                      output_storage[idx][0][:l] = inputs[self.seqs_arg_offset + idx][:l]
                  else:
                      output_storage[idx][0][:] = inputs[self.seqs_arg_offset + idx]
              else:
                  output_storage[idx][0] = inputs[self.seqs_arg_offset + idx].copy()
  
          offset = self.nit_sot_arg_offset + self.n_nit_sot
          other_args = inputs[offset:]
          inner_input_storage = self.fn.input_storage
          nb_mitmot_in = sum(map(len, self.tap_array[: self.n_mit_mot]))
          old_mitmot_input_storage = [None] * nb_mitmot_in
          old_mitmot_input_data = [None] * nb_mitmot_in
          inner_output_storage = self.fn.output_storage
          old_inner_output_storage = [None] * len(inner_output_storage)
          old_inner_output_data = [None] * len(inner_output_storage)
          fn = self.fn.fn
          offset = (
              self.n_seqs
              + sum(map(len, self.tap_array[: self.n_outs]))
              + self.n_shared_outs
          )
          for idx in range(len(other_args)):
              inner_input_storage[idx + offset].storage[0] = other_args[idx]
  
          i = 0
          cond = True
          # ############# THE MAIN LOOP ##############
          # for i in range(n_steps):
          while (i < n_steps) and cond:
              # sequences over which scan iterates
              # 3. collect input slices
              for idx in range(self.n_seqs):
                  if self.vector_seqs[idx]:
                      inner_input_storage[idx].storage[0] = seqs[idx][i : i + 1].reshape(
                          ()
                      )
                  else:
                      inner_input_storage[idx].storage[0] = seqs[idx][i]
  
              offset = self.n_seqs
              for idx in range(self.n_outs):
                  if self.vector_outs[idx]:
                      for tap in self.tap_array[idx]:
                          _idx = (pos[idx] + tap) % store_steps[idx]
                          inner_input_storage[offset].storage[0] = output_storage[idx][0][
                              _idx : _idx + 1
                          ].reshape(())
                          offset += 1
                  else:
                      for tap in self.tap_array[idx]:
                          _idx = (pos[idx] + tap) % store_steps[idx]
                          inner_input_storage[offset].storage[0] = output_storage[idx][0][
                              _idx
                          ]
                          offset += 1
  
              a_offset = self.shared_arg_offset
              o_offset = self.n_outs + self.n_nit_sot
              if i == 0:
                  for j in range(self.n_shared_outs):
                      inner_input_storage[offset].storage[0] = inputs[a_offset + j]
                      offset += 1
              else:
                  for j in range(self.n_shared_outs):
                      inner_input_storage[offset].storage[0] = output_storage[
                          o_offset + j
                      ][0]
                      offset += 1
  
              # 4. collecting slices where the output should be stored
  
              # 4.1. Collect slices for mitmots
              offset = 0
              for idx in range(self.n_mit_mot_outs):
                  if not self.mitmots_preallocated[idx]:
                      inner_output_storage[offset].storage[0] = None
                      offset += 1
  
              # 4.2. Collect slices for mitsots, sitsots and nitsots
              if i != 0:
                  for idx in range(self.n_outs + self.n_nit_sot - self.n_mit_mot):
                      if (
                          store_steps[idx + self.n_mit_mot] == 1
                          or self.vector_outs[idx + self.n_mit_mot]
                      ):
                          inner_output_storage[idx + offset].storage[0] = None
                      else:
                          _pos0 = idx + self.n_mit_mot
                          inner_output_storage[idx + offset].storage[0] = output_storage[
                              _pos0
                          ][0][pos[_pos0]]
              else:
                  for idx in range(self.n_outs + self.n_nit_sot - self.n_mit_mot):
                      inner_output_storage[idx + offset].storage[0] = None
  
              # 4.3. Collect slices for shared outputs
              offset += self.n_outs + self.n_nit_sot - self.n_mit_mot
              for idx in range(self.n_shared_outs):
                  inner_output_storage[idx + offset].storage[0] = None
  
              # 4.4. If there is a condition add it to the mix
              if self.as_while:
                  pdx = offset + self.n_shared_outs
                  inner_output_storage[pdx].storage[0] = None
  
              # 4.5. Keep a reference to the variables (ndarrays, GpuArrays,
              # etc) currently in the output_storage to be able to compare them
              # with the actual outputs of the inner function after its
              # execution. Also keep pointers to their data to be able to detect
              # cases where outputs reused the allocated object but alter the
              # memory region they refer to.
              for idx in range(len(inner_output_storage)):
  
                  var = inner_output_storage[idx].storage[0]
                  old_inner_output_storage[idx] = var
  
                  if var is None:
                      old_inner_output_data[idx] = None
                  elif self.outs_is_tensor[idx]:
                      old_inner_output_data[idx] = var.data
                  else:
                      old_inner_output_data[idx] = var.gpudata
  
              # 4.6. Keep a reference to the variables (ndarrays, GpuArrays,
              # etc) associated with mitmot inputs currently in the
              # input_storage to be able to compare them with the content of the
              # input_storage after the execution of the function. Also keep
              # pointers to their data to be able to detect cases where outputs
              # reused the allocated object but alter the memory region they
              # refer to.
              for idx in range(nb_mitmot_in):
                  var = inner_input_storage[idx + self.n_seqs].storage[0]
                  old_mitmot_input_storage[idx] = var
  
                  if var is None:
                      old_mitmot_input_data[idx] = None
                  elif self.inps_is_tensor[idx + self.n_seqs]:
                      old_mitmot_input_data[idx] = var.data
                  else:
                      old_mitmot_input_data[idx] = var.gpudata
  
              # 5.1 compute outputs
              t0_fn = time.time()
  
              try:
                  fn()
              except Exception:
                  if hasattr(fn, "position_of_error"):
                      # this is a new vm-provided function or c linker
                      # they need this because the exception manipulation
                      # done by raise_with_op is not implemented in C.
                      if hasattr(fn, "thunks"):
                          # For the CVM
                          raise_with_op(
                              self.fn.maker.fgraph,
                              fn.nodes[fn.position_of_error],
                              fn.thunks[fn.position_of_error],
                          )
                      else:
                          # For the c linker
                          # We don't have access from python to all the
                          # temps values So for now, we just don't print
                          # the extra shapes/strides info
                          raise_with_op(
                              self.fn.maker.fgraph, fn.nodes[fn.position_of_error]
                          )
                  else:
                      # old-style linkers raise their own exceptions
                      raise
  
              dt_fn = time.time() - t0_fn
              if self.as_while:
                  pdx = offset + self.n_shared_outs
                  cond = inner_output_storage[pdx].storage[0] == 0
  
              # 5.2. By calling fn() directly instead of calling the aesara
              # function, it is possible that the updates have not been
              # performed. Perform the updates if needed.
              offset_out = len(inner_output_storage) - 1
              if getattr(fn, "need_update_inputs", True):
                  # Update the inputs that have an update function
                  for inp, storage in zip(
                      self.fn.maker.expanded_inputs[::-1], self.fn.input_storage[::-1]
                  ):
                      if inp.update is not None:
                          storage.data = inner_output_storage[offset_out].data
                          offset_out -= 1
  
              t_fn += dt_fn
              offset_out = 0
  
              # 5.3 Copy over the values for mit_mot outputs
              mitmot_inp_offset = 0
              mitmot_out_idx = 0
              for j in range(self.n_mit_mot):
                  for k in self.mit_mot_out_slices[j]:
                      if self.mitmots_preallocated[mitmot_out_idx]:
                          # This output tap has been preallocated.
                          inp_idx = mitmot_inp_offset + self.tap_array[j].index(k)
  
                          # Verify whether the input points to the same data as
                          # it did before the execution of the inner function.
                          old_var = old_mitmot_input_storage[inp_idx]
                          new_var = inner_input_storage[self.n_seqs + inp_idx].storage[0]
                          if old_var is new_var:
                              old_data = old_mitmot_input_data[inp_idx]
                              if self.inps_is_tensor[self.n_seqs + inp_idx]:
                                  same_data = new_var.data == old_data
                              else:
                                  same_data = new_var.gpudata == old_data
                          else:
                              same_data = False
  
                          # If the corresponding input storage still points to
                          # the same data, it has been modified inplace and
                          # nothing needs to be done. Otherwise, recover the
                          # and store it in `outs` as usual
                          if not same_data:
                              output_storage[j][0][k + pos[j]] = inner_input_storage[
                                  self.n_seqs + inp_idx
                              ].storage[0]
  
                      else:
                          # This output tap has not been preallocated, recover
                          # its value as usual
                          output_storage[j][0][k + pos[j]] = inner_output_storage[
                              offset_out
                          ].storage[0]
                          offset_out += 1
  
                      mitmot_out_idx += 1
  
                  mitmot_inp_offset += len(self.tap_array[j])
  
              # 5.4 Copy over the values for mit_sot/sit_sot outputs
              begin = self.n_mit_mot
              end = self.n_outs
              offset_out -= self.n_mit_mot
  
              for j in range(begin, end):
  
                  # Copy the output value to `outs`, if necessary
                  if store_steps[j] == 1 or self.vector_outs[j]:
                      output_storage[j][0][pos[j]] = inner_output_storage[
                          offset_out + j
                      ].storage[0]
                  else:
                      # Check whether the initialization of the output storage
                      # map for this output has been reused.
                      old_var = old_inner_output_storage[offset_out + j]
                      new_var = inner_output_storage[offset_out + j].storage[0]
                      if old_var is new_var:
                          old_data = old_inner_output_data[offset_out + j]
                          if old_data is None:
                              output_reused = False
                          elif self.outs_is_tensor[offset_out + j]:
                              output_reused = new_var.data == old_data
                          else:
                              output_reused = new_var.gpudata == old_data
                      else:
                          output_reused = False
  
                      if not output_reused:
                          try:
                              output_storage[j][0][pos[j]] = inner_output_storage[
                                  offset_out + j
                              ].storage[0]
                          except ValueError as e:
                              if i == 0:
                                  # First iteration, so don't change the
                                  # error message as it can't be the
                                  # case we write about.
                                  raise
                              ne = ValueError(
                                  "An output of the scan has changed shape. "
                                  "This may be caused by a pushout optimization."
                                  " Try adding "
                                  "'optimizer_excluding=scanOp_pushout_output' "
                                  "to your Aesara flags."
                              )
                              raise ne from e
  
              # 5.5 Copy over the values for nit_sot outputs
              begin = end
              end += self.n_nit_sot
              for j in range(begin, end):
  
                  if i == 0:
                      jout = j + offset_out
                      shape = (store_steps[j],) + inner_output_storage[jout].storage[
                          0
                      ].shape
                      dtype = inner_output_storage[jout].storage[0].dtype
                      if (
                          output_storage[j][0] is None
                          or output_storage[j][0].shape[0] < store_steps[j]
                          or output_storage[j][0].shape[1:] != shape[1:]
                          or output_storage[j][0].dtype != dtype
                      ):
                          output_storage[j][0] = node.outputs[j].type.value_zeros(shape)
                      elif output_storage[j][0].shape[0] != store_steps[j]:
                          output_storage[j][0] = output_storage[j][0][: store_steps[j]]
                      output_storage[j][0][pos[j]] = inner_output_storage[jout].storage[0]
                  elif store_steps[j] == 1 or self.vector_outs[j]:
                      output_storage[j][0][pos[j]] = inner_output_storage[
                          j + offset_out
                      ].storage[0]
                  else:
                      # Check whether the initialization of the output storage map
                      # for this output has been reused.
                      old_var = old_inner_output_storage[offset_out + j]
                      old_data = old_inner_output_data[offset_out + j]
                      new_var = inner_output_storage[offset_out + j].storage[0]
                      if old_var is new_var:
                          if old_data is None:
                              output_reused = False
                          elif self.outs_is_tensor[offset_out + j]:
                              output_reused = new_var.data == old_data
                          else:
                              output_reused = new_var.gpudata == old_data
                      else:
                          output_reused = False
  
                      if not output_reused:
                          output_storage[j][0][pos[j]] = inner_output_storage[
                              j + offset_out
                          ].storage[0]
  
              # 5.6 Copy over the values for outputs corresponding to shared
              # variables
              begin = end
              end += self.n_shared_outs
              for j in range(begin, end):
                  jout = j + offset_out
                  output_storage[j][0] = inner_output_storage[jout].storage[0]
  
              pos = [(idx + 1) % store for idx, store in zip(pos, store_steps)]
              i = i + 1
  
          # 6. Check if you need to re-order output buffers
          begin = self.n_mit_mot
          end = self.n_outs + self.n_nit_sot
          for idx in range(begin, end):
              if store_steps[idx] < i - self.mintaps[idx] and pos[idx] < store_steps[idx]:
  
                  pdx = pos[idx]
                  if pdx >= store_steps[idx] // 2:
                      # It seems inefficient to copy the bigger part of the
                      # array over, and back, but it is the only way that
                      # there is no overlap in the areas of out[idx][0] that
                      # are read and written.
                      # This way, there will be no information overwritten
                      # before it is read (as it used to happen).
                      shape = (pdx,) + output_storage[idx][0].shape[1:]
                      tmp = node.outputs[idx].type.value_zeros(shape)
                      tmp[:] = output_storage[idx][0][:pdx]
                      output_storage[idx][0][: store_steps[idx] - pdx] = output_storage[
                          idx
                      ][0][pdx:]
                      output_storage[idx][0][store_steps[idx] - pdx :] = tmp
                      del tmp
                  else:
                      shape = (store_steps[idx] - pdx,) + output_storage[idx][0].shape[1:]
                      tmp = node.outputs[idx].type.value_zeros(shape)
                      tmp[:] = output_storage[idx][0][pdx:]
                      output_storage[idx][0][store_steps[idx] - pdx :] = output_storage[
                          idx
                      ][0][:pdx]
                      output_storage[idx][0][: store_steps[idx] - pdx] = tmp
                      del tmp
              # This would normally happen only when doing truncated
              # backpropagation through time. In such a scenario Scan is
              # expected to return 0 for all entries for which the gradient is
              # not actually computed
              elif store_steps[idx] > i - self.mintaps[idx]:
                  output_storage[idx][0][i - self.mintaps[idx] :] = 0
                  # This is a fix for a bug introduced by while. If you say
                  # you want to loop up to a condition, you expect the output
                  # to have that length ( and not the maximal length possible)
                  #
                  # Without this the behaviour of a scan op is not consistent
                  # if optimization gets applied compared to when optimization
                  # do not get applied
                  if i < n_steps:
                      # The reason I don't use out[idx][0][:i] is because for
                      # certain outputs (those with multiple taps),
                      # outs[idx][0] has more than n_steps entries, with the
                      # initial state  at the beginning. When indexing in it I
                      # usually have to do something like
                      # outs[idx][0][i+offset]. To do something similar here,
                      # I would have first to compute the maximal tap for
                      # every output and then do outs[0][:i+maximal_tap],
                      # which implies I think more computations then this
                      # little trick that I used
                      output_storage[idx][0] = output_storage[idx][0][: -(n_steps - i)]
  
          # We never reuse the input or output storage of the
          # inner function so we clear it.
          for i_s in inner_input_storage:
              i_s.storage[0] = None
          for o_s in inner_output_storage:
              o_s.storage[0] = None
  
          t_call = time.time() - t0_call
          # NOTE: make this match what's in function.types.Function
          # and this little string helps us to find this spot:
          # "PROFILE_CODE"
  
          if hasattr(self.fn.maker, "profile") and self.fn.maker.profile:
              profile = self.fn.maker.profile
              profile.callcount += 1
              profile.nbsteps += n_steps
              profile.call_time += t_call
              profile.vm_call_time += t_fn
              if hasattr(self.fn.fn, "update_profile"):
                  self.fn.fn.update_profile(profile)
  
          self.t_call = t_call
          self.t_fn = t_fn
  
      def infer_shape(self, fgraph, node, input_shapes):
          # input_shapes correspond to the shapes of node.inputs
          for inp, inp_shp in zip(node.inputs, input_shapes):
              assert inp_shp is None or len(inp_shp) == inp.type.ndim
  
          # Here we build 2 variables;
          # - A list `inner_ins_shapes`, such that inner_ins_shapes[i] is the
          #   shape of self.inputs[i]
          # - A dictionary `out_equivalent` containing, for every inner input,
          #   an equivalent variable computed from the outer inputs.
          #   NOTE : For non-sequences, this equivalence is trivial. For
          #   sequences and recurrent states, there is no direct equivalence
          #   between outer and inner inputs. However, because every iteration
          #   of the Scan needs to give the same output shapes, we can give an
          #   equivalence between these inner inputs and the subelements of the
          #   corresponding outer inputs that the Scan would use as input for
          #   any given iteration. For simplicity, we use iteration 0.
          inner_ins_shapes = []
          out_equivalent = OrderedDict()
  
          # The two following blocks are commented as it cause in some
          # cases extra scans in the graph. See gh-XXX for the
          # investigation.
  
          # We skip the first outer input as it is the total or current number
          # of iterations.
          # sequences
          seqs_shape = [x[1:] for x in input_shapes[1 : 1 + self.n_seqs]]
          # We disable extra infer_shape for now. See gh-3765.
          extra_infer_shape = False
  
          if extra_infer_shape:
              inner_seqs = self.inputs[: self.n_seqs]
              outer_seqs = node.inputs[1 : 1 + self.n_seqs]
              for in_s, out_s in zip(inner_seqs, outer_seqs):
                  out_equivalent[in_s] = out_s[0]
  
              # mit_mot, mit_sot, sit_sot
              outer_inp_idx = 1 + self.n_seqs
              inner_inp_idx = self.n_seqs
          else:
              outer_inp_idx = 0
          n_outs = self.n_mit_mot + self.n_mit_sot + self.n_sit_sot
          outs_shape = []
          for idx in range(n_outs):
              mintap = abs(min(self.tap_array[idx]))
              for k in self.tap_array[idx]:
                  outs_shape += [input_shapes[idx + self.n_seqs + 1][1:]]
                  if extra_infer_shape:
                      corresponding_tap = node.inputs[outer_inp_idx][mintap + k]
                      out_equivalent[self.inputs[inner_inp_idx]] = corresponding_tap
                      inner_inp_idx += 1
              outer_inp_idx += 1
  
          # shared_outs
          offset = 1 + self.n_seqs + n_outs
          for idx in range(self.n_shared_outs):
              outs_shape += [input_shapes[idx + offset]]
  
          # non_sequences
          offset += self.n_nit_sot + self.n_shared_outs
          inner_ins_shapes = seqs_shape + outs_shape + input_shapes[offset:]
          assert len(inner_ins_shapes) == len(self.inputs)
  
          # Non-sequences have a direct equivalent from self.inputs in
          # node.inputs
          inner_non_sequences = self.inputs[len(seqs_shape) + len(outs_shape) :]
          for in_ns, out_ns in zip(inner_non_sequences, node.inputs[offset:]):
              out_equivalent[in_ns] = out_ns
  
          if self.as_while:
              self_outs = self.outputs[:-1]
          else:
              self_outs = self.outputs
          outs_shape = infer_shape(
              outs=self_outs, inputs=self.inputs, input_shapes=inner_ins_shapes
          )
          # Will be used to check if outs_shape can be expressed without using
          # variables in self.inputs.
          # The shapes of node.inputs are valid.
          validator = Validator(
              valid=input_shapes, invalid=self.inputs, valid_equivalent=out_equivalent
          )
  
          offset = 1 + self.n_seqs
          scan_outs = [x for x in input_shapes[offset : offset + n_outs]]
          offset += n_outs
          outs_shape_n = self.n_mit_mot_outs + self.n_mit_sot + self.n_sit_sot
          for x in range(self.n_nit_sot):
              out_shape_x = outs_shape[outs_shape_n + x]
              if out_shape_x is None:
                  # This output is not a tensor, and has no shape
                  scan_outs.append(None)
              else:
                  # We need to make sure that we can compute the shapes from
                  # node.inputs, and constants, without using the variables
                  # in the inner function.
                  r = node.outputs[n_outs + x]
                  assert r.ndim == 1 + len(out_shape_x)
                  shp = [node.inputs[offset + self.n_shared_outs + x]]
                  for i, shp_i in zip(range(1, r.ndim), out_shape_x):
                      # Validate shp_i. v_shape_i is either None (if invalid),
                      # or a (variable, Boolean) tuple. The Boolean indicates
                      # whether variable is shp_i (if True), or an valid
                      # equivalent (if False). Here, we only need the variable.
                      v_shp_i = validator.check(shp_i)
                      if v_shp_i is None:
                          if hasattr(r, "broadcastable") and r.broadcastable[i]:
                              shp.append(1)
                          else:
                              shp.append(Shape_i(i)(r))
                      else:
                          # It can (or at least, an equivalent variable can)
                          shp.append(v_shp_i[0])
                  scan_outs.append(tuple(shp))
  
          scan_outs += [x for x in input_shapes[offset : offset + self.n_shared_outs]]
          # if we are dealing with a repeat-until, then we do not know the
          # leading dimension so we replace it for every entry with Shape_i
          if self.as_while:
              scan_outs_init = scan_outs
              scan_outs = []
              for o, x in zip(node.outputs, scan_outs_init):
                  if x is None:
                      scan_outs.append(None)
                  else:
                      scan_outs.append((Shape_i(0)(o),) + x[1:])
          return scan_outs
  
      def connection_pattern(self, node):
  
          # We cache the result of this function because, with a previous
          # implementation that repeatedly called grad, there were cases
          # where calls to aesara.grad() took as much as 4h for functions
          # containing many nested scans.
          if hasattr(node.tag, "connection_pattern"):
              return node.tag.connection_pattern
  
          # Obtain the connection pattern of the inner function.
          inner_connect_pattern = io_connection_pattern(self.inputs, self.outputs)
  
          # Initially assume no outer input is connected to any outer output
          connection_pattern = [[False for output in node.outputs] for x in node.inputs]
  
          # For every possible pair of outer input and outer output, iterate
          # over every possible pairing of their corresponding inner inputs
          # and inner outputs and, if one such pair of inner variables is
          # connected than the pair of outer variables is connected.
          for outer_oidx in range(len(node.outputs)):
              inner_oidxs = self.var_mappings["inner_out_from_outer_out"][outer_oidx]
  
              for outer_iidx in range(len(node.inputs)):
                  inner_iidxs = self.var_mappings["inner_inp_from_outer_inp"][outer_iidx]
  
                  for inner_oidx in inner_oidxs:
                      for inner_iidx in inner_iidxs:
  
                          if inner_connect_pattern[inner_iidx][inner_oidx]:
                              connection_pattern[outer_iidx][outer_oidx] = True
                              break
  
                      if connection_pattern[outer_iidx][outer_oidx]:
                          break
  
          # Applying Floyd-Warshall to find all paths connecting inputs to
          # outputs. Note that if `x` is an input to `y_t` and `y_tm1` is an
          # input to `z_t` then `x` is an input to `z_t`.
  
          n_outs = len(node.outputs)
  
          for steps in range(n_outs):
              for iidx in range(n_outs):
                  for jidx in range(n_outs):
  
                      # Get the idx of the outer input corresponding to that
                      # outer output
                      j_inp_idx = self.var_mappings["outer_inp_from_outer_out"][jidx]
  
                      if j_inp_idx != -1:
                          if connection_pattern[j_inp_idx][iidx] is True:
                              for k in range(len(connection_pattern)):
                                  if connection_pattern[k][jidx]:
                                      connection_pattern[k][iidx] = True
  
          node.tag.connection_pattern = connection_pattern
          return connection_pattern
  
      def get_oinp_iinp_iout_oout_mappings(self):
          """
          Compute and return dictionary mappings between the inputs and
          outputs of the inner function and the inputs and outputs of the Scan
          node in the outer graph.
  
          The return value is a dictionary in which the keys are the names of
          the individual mappings and the values are the mapping dictionaries
          themselves. In dictionaries representing mappings to outer variables,
          the values are individual integer indices. In dictionaries
          representing mappings to inner variables, the values are sequences of
          indices because multiple inner variables can be associated with the
          same state.
  
          """
          # Lists for outer variables contain individual indices, lists for
          # inner variables contain sequences of indices because many inner
          # variables can be associated with the same outer variable. The list
          # and indices are initialized already containing the data associated
          # with the timestep index, the first outer input.
          outer_input_indices = [0]
          inner_input_indices = [[]]
          inner_output_indices = [[]]
          outer_output_indices = [-1]
  
          outer_iidx = 1
          inner_iidx = 0
          inner_oidx = 0
          outer_oidx = 0
  
          # Handle sequences inputs
          for i in range(self.info["n_seqs"]):
              outer_input_indices.append(outer_iidx)
              inner_input_indices.append([inner_iidx])
              inner_output_indices.append([])
              outer_output_indices.append(-1)
  
              outer_iidx += 1
              inner_iidx += 1
              inner_oidx += 0
              outer_oidx += 0
  
          # Handle mitmots, mitsots and sitsots variables
          for i in range(len(self.info["tap_array"])):
              nb_input_taps = len(self.info["tap_array"][i])
  
              if i < self.n_mit_mot:
                  nb_output_taps = len(self.mit_mot_out_slices[i])
              else:
                  nb_output_taps = 1
  
              outer_input_indices.append(outer_iidx)
              inner_input_indices.append(
                  list(range(inner_iidx, inner_iidx + nb_input_taps))
              )
              inner_output_indices.append(
                  list(range(inner_oidx, inner_oidx + nb_output_taps))
              )
              outer_output_indices.append(outer_oidx)
  
              outer_iidx += 1
              inner_iidx += nb_input_taps
              inner_oidx += nb_output_taps
              outer_oidx += 1
  
          # This is needed because, for outer inputs (and for outer inputs only)
          # nitsots come *after* shared variables.
          outer_iidx += self.info["n_shared_outs"]
  
          # Handle nitsots variables
          for i in range(self.n_nit_sot):
              outer_input_indices.append(outer_iidx)
              inner_input_indices.append([])
              inner_output_indices.append([inner_oidx])
              outer_output_indices.append(outer_oidx)
  
              outer_iidx += 1
              inner_iidx += 0
              inner_oidx += 1
              outer_oidx += 1
  
          # This is needed because, for outer inputs (and for outer inputs only)
          # nitsots come *after* shared variables.
          outer_iidx -= self.info["n_shared_outs"] + self.n_nit_sot
  
          # Handle shared states
          for i in range(self.info["n_shared_outs"]):
              outer_input_indices.append(outer_iidx)
              inner_input_indices.append([inner_iidx])
              inner_output_indices.append([inner_oidx])
              outer_output_indices.append(outer_oidx)
  
              outer_iidx += 1
              inner_iidx += 1
              inner_oidx += 1
              outer_oidx += 1
  
          # This is needed because, for outer inputs (and for outer inputs only)
          # nitsots come *after* shared variables.
          outer_iidx += self.n_nit_sot
  
          # Handle non-sequence inputs
          # Note : the number of non-sequence inputs is not stored in self.info
          # so it has to be inferred from the number of inner inputs that remain
          # to be handled
          for i in range(len(self.inputs) - inner_iidx):
              outer_input_indices.append(outer_iidx)
              inner_input_indices.append([inner_iidx])
              inner_output_indices.append([])
              outer_output_indices.append(-1)
  
              outer_iidx += 1
              inner_iidx += 1
              inner_oidx += 0
              outer_oidx += 0
  
          # With the global mapping inferred, the individual mappings
          # can be produced
          mappings = {
              "outer_inp_from_outer_out": {},
              "inner_inp_from_outer_out": {},
              "inner_out_from_outer_out": {},
              "inner_inp_from_outer_inp": {},
              "inner_out_from_outer_inp": {},
              "outer_out_from_outer_inp": {},
              "outer_inp_from_inner_inp": {},
              "inner_out_from_inner_inp": {},
              "outer_out_from_inner_inp": {},
              "outer_inp_from_inner_out": {},
              "inner_inp_from_inner_out": {},
              "outer_out_from_inner_out": {},
          }
  
          for (oinp, iinp, iout, oout) in zip(
              outer_input_indices,
              inner_input_indices,
              inner_output_indices,
              outer_output_indices,
          ):
  
              if oout != -1:
                  mappings["outer_inp_from_outer_out"][oout] = oinp
                  mappings["inner_inp_from_outer_out"][oout] = iinp
                  mappings["inner_out_from_outer_out"][oout] = iout
  
              if oinp != -1:
                  mappings["inner_inp_from_outer_inp"][oinp] = iinp
                  mappings["inner_out_from_outer_inp"][oinp] = iout
                  mappings["outer_out_from_outer_inp"][oinp] = oout
  
              for idx in iinp:
                  mappings["outer_inp_from_inner_inp"][idx] = oinp
                  mappings["inner_out_from_inner_inp"][idx] = iout
                  mappings["outer_out_from_inner_inp"][idx] = oout
  
              for idx in iout:
                  mappings["outer_inp_from_inner_out"][idx] = oinp
                  mappings["inner_inp_from_inner_out"][idx] = iinp
                  mappings["outer_out_from_inner_out"][idx] = oout
  
          return mappings
  
      def L_op(self, inputs, outs, dC_douts):
          if not isinstance(outs, (list, tuple)):
              outs = [outs]
          # `grad_step` equals the number of steps the original scan node has
          # done (if the original scan is a while loop than this number is the
          # length of the output sequence)
          # We do not know what kind of outputs the original scan has, so we
          # try first to see if it has a nit_sot output, then a sit_sot and
          # then a mit_sot
          if self.n_nit_sot > 0:
              grad_steps = self.outer_nitsot_outs(outs)[0].shape[0]
          elif self.n_sit_sot > 0:
              grad_steps = self.outer_sitsot_outs(outs)[0].shape[0] - 1
          elif self.n_mit_sot > 0:
              grad_steps = (
                  self.outer_mitsot_outs(outs)[0].shape[0] + self.mintaps[self.n_mit_mot]
              )
          else:
              grad_steps = inputs[0]
          if self.as_while:
              n_steps = outs[0].shape[0]
  
          # Restrict the number of grad steps according to
          # self.truncate_gradient
          if self.truncate_gradient != -1:
              grad_steps = minimum(grad_steps, self.truncate_gradient)
  
          self_inputs = self.inputs
          self_outputs = self.outputs
          # differentiable inputs
          diff_inputs = (
              self.inner_seqs(self_inputs)
              + self.inner_mitmot(self_inputs)
              + self.inner_mitsot(self_inputs)
              + self.inner_sitsot(self_inputs)
              + self.inner_non_seqs(self_inputs)
          )
          diff_outputs = (
              self.inner_mitmot_outs(self_outputs)
              + self.inner_mitsot_outs(self_outputs)
              + self.inner_sitsot_outs(self_outputs)
              + self.inner_nitsot_outs(self_outputs)
          )
          scan_node = outs[0].owner
          connection_pattern = self.connection_pattern(scan_node)
  
          def get_inp_idx(iidx):
              if iidx < self.n_seqs:
                  return 1 + iidx
              oidx = 1 + self.n_seqs
              iidx = iidx - self.n_seqs
              for taps in self.mitmot_taps():
                  if len(taps) > iidx:
                      return oidx
                  else:
                      oidx += 1
                      iidx -= len(taps)
              for taps in self.mitsot_taps():
                  if len(taps) > iidx:
                      return oidx
                  else:
                      oidx += 1
                      iidx -= len(taps)
  
              if iidx < self.info["n_sit_sot"]:
                  return oidx + iidx
              else:
                  return oidx + iidx + self.info["n_nit_sot"]
  
          def get_out_idx(iidx):
              oidx = 0
              for taps in self.mitmot_out_taps():
                  if len(taps) > iidx:
                      return oidx
                  else:
                      oidx += 1
                      iidx -= len(taps)
              return oidx + iidx
  
          def compute_all_gradients(known_grads):
              y_s = known_grads.keys()
              g_y_s = known_grads.values()
  
              for g_y in g_y_s:
                  if str(g_y.dtype) in integer_dtypes:
                      raise TypeError(
                          "Gradients may never be integers but g_y "
                          "has type " + str(g_y.type)
                      )
  
              out_indices = [get_out_idx(self_outputs.index(y)) for y in y_s]
  
              connected_inputs = [
                  i
                  for i in range(len(scan_node.inputs))
                  if any([connection_pattern[i][odx] for odx in out_indices])
              ]
  
              wrt = [
                  x
                  for x in graph_inputs(y_s)
                  if (x in diff_inputs)
                  and get_inp_idx(self_inputs.index(x)) in connected_inputs
              ]
              gmp = OrderedDict()
  
              # Required in case there is a pair of variables X and Y, with X
              # used to compute Y, for both of which there is an external
              # gradient signal. Without this, the total gradient signal on X
              # will be the external gradient  signalknown_grads[X]. With this,
              # it will be the sum of the external gradient signal and the
              # gradient obtained by propagating Y's external gradient signal
              # to X.
              known_grads = OrderedDict([(k.copy(), v) for (k, v) in known_grads.items()])
  
              grads = grad(
                  cost=None,
                  known_grads=known_grads,
                  wrt=wrt,
                  consider_constant=wrt,
                  disconnected_inputs="ignore",
                  return_disconnected="None",
                  null_gradients="return",
              )
  
              for i in range(len(wrt)):
                  gmp[wrt[i]] = grads[i]
  
              rval = [gmp.get(p, None) for p in diff_inputs]
              return rval
  
          dC_dinps_t = [None for inp in diff_inputs]
          disconnected_dC_dinps_t = [True for inp in diff_inputs]
          dC_dXts = []
          Xts = []
          for idx, Xt in enumerate(diff_outputs):
  
              # We are looking for x[t-1] for a given x[t]
              if idx >= self.n_mit_mot_outs:
                  Xt_placeholder = safe_new(Xt)
                  Xts.append(Xt_placeholder)
  
              # Different processing based on whether Xt is a nitsot output
              # or not. NOTE : This cannot be done by using
              # "if Xt not in self.inner_nitsot_outs(self_outputs)" because
              # the exact same variable can be used as multiple outputs.
              idx_nitsot_start = (
                  self.info["n_mit_mot"] + self.info["n_mit_sot"] + self.info["n_sit_sot"]
              )
              idx_nitsot_end = idx_nitsot_start + self.info["n_nit_sot"]
              if idx < idx_nitsot_start or idx >= idx_nitsot_end:
                  # What we do here is loop through dC_douts and collect all
                  # those that are connected to the specific one and do an
                  # upcast on all of their dtypes to get the dtype for this
                  # specific output. Deciding if the gradient with this
                  # specific previous step is defined or not is done somewhere
                  # else.
                  dtypes = []
                  states = (
                      self.inner_mitmot(self_inputs)
                      + self.inner_mitsot(self_inputs)
                      + self.inner_sitsot(self_inputs)
                  )
  
                  for pos, inp in enumerate(states):
                      if inp in graph_inputs([Xt]):
                          # Get the index of the outer output that to which
                          # the state variable 'inp' corresponds.
                          outer_oidx = self.var_mappings["outer_out_from_inner_inp"][
                              self.n_seqs + pos
                          ]
  
                          if not isinstance(dC_douts[outer_oidx].type, DisconnectedType):
                              dtypes.append(dC_douts[outer_oidx].dtype)
                  if dtypes:
                      new_dtype = aesara.scalar.upcast(*dtypes)
                  else:
                      new_dtype = config.floatX
                  dC_dXt = safe_new(Xt, dtype=new_dtype)
              else:
                  if isinstance(dC_douts[idx].type, DisconnectedType):
                      continue
                  dC_dXt = safe_new(dC_douts[idx][0])
              dC_dXts.append(dC_dXt)
  
          known_grads = OrderedDict()
          dc_dxts_idx = 0
          for i in range(len(diff_outputs)):
              if i < idx_nitsot_start or i >= idx_nitsot_end:
                  if diff_outputs[i] in known_grads:
                      known_grads[diff_outputs[i]] += dC_dXts[dc_dxts_idx]
                  else:
                      known_grads[diff_outputs[i]] = dC_dXts[dc_dxts_idx]
                  dc_dxts_idx += 1
              else:
                  if isinstance(dC_douts[i].type, DisconnectedType):
                      continue
                  else:
                      if diff_outputs[i] in known_grads:
                          known_grads[diff_outputs[i]] += dC_dXts[dc_dxts_idx]
                      else:
                          known_grads[diff_outputs[i]] = dC_dXts[dc_dxts_idx]
                      dc_dxts_idx += 1
          dC_dinps_t = compute_all_gradients(known_grads)
  
          # mask inputs that get no gradients
          for dx in range(len(dC_dinps_t)):
              if not dC_dinps_t[dx]:
                  dC_dinps_t[dx] = aet.zeros_like(diff_inputs[dx])
              else:
                  disconnected_dC_dinps_t[dx] = False
                  for Xt, Xt_placeholder in zip(diff_outputs[self.n_mit_mot_outs :], Xts):
                      tmp = forced_replace(dC_dinps_t[dx], Xt, Xt_placeholder)
                      dC_dinps_t[dx] = tmp
  
          # construct dX_dtm1
          dC_dXtm1s = []
          for pos, x in enumerate(dC_dinps_t[self.n_seqs :]):
  
              # Get the index of the first inner input corresponding to the
              # pos-ieth inner input state
              idxs = self.var_mappings["inner_out_from_inner_inp"][self.n_seqs + pos]
  
              # Check if the pos-th input is associated with one of the
              # recurrent states
              x_is_state = pos < sum([len(t) for t in self.tap_array])
  
              if x_is_state and len(idxs) > 0:
                  opos = idxs[0]
                  dC_dXtm1s.append(safe_new(dC_dXts[opos]))
                  if hasattr(x, "dtype") and x.dtype != dC_dXts[opos].dtype:
                      dC_dinps_t[pos + self.n_seqs] = x.astype(dC_dXts[opos].dtype)
              else:
                  dC_dXtm1s.append(safe_new(x))
  
          for dx, dC_dXtm1 in enumerate(dC_dXtm1s):
              if isinstance(dC_dinps_t[dx + self.n_seqs].type, NullType):
                  # The accumulated gradient is undefined
                  pass
              elif isinstance(dC_dXtm1.type, NullType):
                  # The new gradient is undefined, this makes the accumulated
                  # gradient undefined as weell
                  dC_dinps_t[dx + self.n_seqs] = dC_dXtm1
              else:
                  dC_dinps_t[dx + self.n_seqs] += dC_dXtm1
          # Construct scan op
          # Seqs
          if self.as_while:
              # equivalent to x[:n_steps][::-1]
              outer_inp_seqs = [x[n_steps - 1 :: -1] for x in inputs[1 : 1 + self.n_seqs]]
          else:
              outer_inp_seqs = [x[::-1] for x in inputs[1 : 1 + self.n_seqs]]
          for idx in range(self.n_mit_mot + self.n_mit_sot):
              mintap = np.min(self.tap_array[idx])
              if idx < self.n_mit_mot:
                  outmaxtap = np.max(self.mitmot_out_taps()[idx])
              else:
                  outmaxtap = 0
              seq = outs[idx]
              for k in self.tap_array[idx]:
                  if outmaxtap - k != 0:
                      nw_seq = seq[k - mintap : -(outmaxtap - k)][::-1]
                  else:
                      nw_seq = seq[k - mintap :][::-1]
                  outer_inp_seqs.append(nw_seq)
          outer_inp_seqs += [x[:-1][::-1] for x in self.outer_sitsot_outs(outs)]
          for x in self.outer_nitsot_outs(dC_douts):
              if not isinstance(x.type, DisconnectedType):
                  if self.as_while:
                      # equivalent to x[:n_steps][::-1]
                      outer_inp_seqs.append(x[n_steps - 1 :: -1])
                  else:
                      outer_inp_seqs.append(x[::-1])
  
          if hasattr(inputs[0].tag, "test_value"):
              # Here we tests that the new scan input sequence all have
              # the same shape[0]. This is a properties that the scan()
              # fct add and we want to keep it for all Scan op.  This is
              # used in T_Scan.test_grad_multiple_outs_taps to test
              # that.
              if self.as_while:
                  n = n_steps.tag.test_value
              else:
                  n = inputs[0].tag.test_value
              for taps, x in zip(self.mitsot_taps(), self.outer_mitsot_outs(outs)):
                  mintap = np.min(taps)
                  if hasattr(x[::-1][:mintap], "test_value"):
                      assert x[::-1][:mintap].tag.test_value.shape[0] == n
              for x in self.outer_sitsot_outs(outs):
                  if hasattr(x[::-1][:-1].tag, "test_value"):
                      assert x[::-1][:-1].tag.test_value.shape[0] == n
              for x in self.outer_nitsot_outs(outs):
                  if hasattr(x[::-1].tag, "test_value"):
                      if self.as_while:
                          assert x[n_steps - 1 :: -1].tag.test_value.shape[0] == n
                      else:
                          assert x[::-1].tag.test_value.shape[0] == n
          outer_inp_seqs += [
              x[::-1][: np.min(taps)]
              for taps, x in zip(self.mitsot_taps(), self.outer_mitsot_outs(outs))
          ]
          outer_inp_seqs += [x[::-1][:-1] for x in self.outer_sitsot_outs(outs)]
          outer_inp_seqs += [x[::-1] for x in self.outer_nitsot_outs(outs)]
  
          # Restrict the length of the outer sequences to the number of grad
          # steps
          outer_inp_seqs = [s_[:grad_steps] for s_ in outer_inp_seqs]
  
          inner_inp_seqs = self.inner_seqs(self_inputs)
          inner_inp_seqs += self.inner_mitmot(self_inputs)
          inner_inp_seqs += self.inner_mitsot(self_inputs)
          inner_inp_seqs += self.inner_sitsot(self_inputs)
          inner_inp_seqs += self.inner_nitsot_outs(dC_dXts)
          inner_inp_seqs += Xts
          # mitmot
          outer_inp_mitmot = []
          inner_inp_mitmot = []
          inner_out_mitmot = []
          mitmot_inp_taps = []
          mitmot_out_taps = []
          type_outs = []
          out_pos = 0
          ins_pos = self.n_seqs
          n_mitmot_outs = 0
          n_mitmot_inps = 0
  
          for idx in range(self.n_mit_mot):
              if isinstance(dC_douts[idx].type, DisconnectedType):
                  out = outs[idx]
                  outer_inp_mitmot.append(aet.zeros_like(out))
              else:
                  outer_inp_mitmot.append(dC_douts[idx][::-1])
              mitmot_inp_taps.append([])
              mitmot_out_taps.append([])
              undefined_msg = None
              through_shared = False
              disconnected = True
  
              for jdx in range(len(self.mit_mot_out_slices[idx])):
                  inner_inp_mitmot.append(dC_dXts[out_pos])
                  mitmot_inp_taps[idx].append(-self.mit_mot_out_slices[idx][jdx])
                  n_mitmot_inps += 1
                  out_pos += 1
  
              for jdx in range(len(self.tap_array[idx])):
                  tap = -self.tap_array[idx][jdx]
  
                  # Only create a new inner input if there is not already one
                  # associated with this input tap
                  if tap not in mitmot_inp_taps[idx]:
                      inner_inp_mitmot.append(dC_dXtm1s[ins_pos - self.n_seqs])
  
                  if isinstance(dC_dinps_t[ins_pos].type, NullType):
                      # We cannot use Null in the inner graph, so we
                      # use a zero tensor of the appropriate shape instead.
                      inner_out_mitmot.append(
                          aet.zeros(diff_inputs[ins_pos].shape, dtype=config.floatX)
                      )
                      undefined_msg = dC_dinps_t[ins_pos].type.why_null
                  else:
                      new_inner_out_mitmot = dC_dinps_t[ins_pos]
  
                      # If there is already an inner input associated with that
                      # input tap, make sure the computation of the new output
                      # uses it instead of the input it's currently using
                      if tap in mitmot_inp_taps[idx]:
                          to_replace = dC_dXtm1s[ins_pos - self.n_seqs]
                          replacement_idx = len(mitmot_inp_taps[idx]) - mitmot_inp_taps[
                              idx
                          ].index(tap)
                          replacement = inner_inp_mitmot[-replacement_idx]
  
                          self.tap_array[idx]
                          new_inner_out_mitmot = clone_replace(
                              new_inner_out_mitmot, replace=[(to_replace, replacement)]
                          )
  
                      inner_out_mitmot.append(new_inner_out_mitmot)
  
                  if not disconnected_dC_dinps_t[ins_pos]:
                      disconnected = False
  
                  for _sh in self.inner_shared(self_inputs):
                      if _sh in graph_inputs([dC_dinps_t[ins_pos]]):
                          through_shared = True
  
                  ins_pos += 1
                  n_mitmot_outs += 1
                  mitmot_out_taps[idx].append(-self.tap_array[idx][jdx])
  
                  # Only add the tap as a new input tap if needed
                  if tap not in mitmot_inp_taps[idx]:
                      n_mitmot_inps += 1
                      mitmot_inp_taps[idx].append(-self.tap_array[idx][jdx])
  
              if undefined_msg:
                  type_outs.append(undefined_msg)
              elif through_shared:
                  type_outs.append("through_shared")
              elif disconnected:
                  type_outs.append("disconnected")
              else:
                  type_outs.append("connected")
  
          offset = self.n_mit_mot
          for idx in range(self.n_mit_sot):
              if isinstance(dC_douts[idx + offset].type, DisconnectedType):
                  outer_inp_mitmot.append(outs[idx + offset].zeros_like())
              else:
                  outer_inp_mitmot.append(dC_douts[idx + offset][::-1])
              mitmot_inp_taps.append([])
              mitmot_out_taps.append([])
              idx_tap = idx + self.n_mit_mot
              inner_inp_mitmot.append(dC_dXts[out_pos])
              out_pos += 1
              n_mitmot_inps += 1
              undefined_msg = None
              through_shared = False
              disconnected = True
              mitmot_inp_taps[idx + offset].append(0)
              for jdx in range(len(self.tap_array[idx_tap])):
                  inner_inp_mitmot.append(dC_dXtm1s[ins_pos - self.n_seqs])
  
                  if isinstance(dC_dinps_t[ins_pos].type, NullType):
                      # We cannot use Null in the inner graph, so we
                      # use a zero tensor of the appropriate shape instead.
                      inner_out_mitmot.append(
                          aet.zeros(diff_inputs[ins_pos].shape, dtype=config.floatX)
                      )
                      undefined_msg = dC_dinps_t[ins_pos].type.why_null
                  else:
                      inner_out_mitmot.append(dC_dinps_t[ins_pos])
  
                  mitmot_inp_taps[idx + offset].append(-self.tap_array[idx_tap][jdx])
                  mitmot_out_taps[idx].append(-self.tap_array[idx_tap][jdx])
                  if not disconnected_dC_dinps_t[ins_pos]:
                      disconnected = False
                  for _sh in self.inner_shared(self_inputs):
                      if _sh in graph_inputs([dC_dinps_t[ins_pos]]):
                          through_shared = True
  
                  n_mitmot_inps += 1
                  ins_pos += 1
                  n_mitmot_outs += 1
  
              if undefined_msg:
                  type_outs.append(undefined_msg)
              elif through_shared:
                  type_outs.append("through_shared")
              elif disconnected:
                  type_outs.append("disconnected")
              else:
                  type_outs.append("connected")
  
          offset += self.n_mit_sot
          for idx in range(self.n_sit_sot):
              mitmot_inp_taps.append([0, 1])
              mitmot_out_taps.append([1])
              through_shared = False
              if not isinstance(dC_douts[idx + offset].type, DisconnectedType):
                  outer_inp_mitmot.append(dC_douts[idx + offset][::-1])
              else:
                  if isinstance(dC_dinps_t[ins_pos].type, NullType):
                      # Cannot use dC_dinps_t[ins_pos].dtype, so we use
                      # floatX instead, as it is a dummy value that will not
                      # be used anyway.
                      outer_inp_mitmot.append(
                          aet.zeros(outs[idx + offset].shape, dtype=config.floatX)
                      )
                  else:
                      outer_inp_mitmot.append(
                          aet.zeros(
                              outs[idx + offset].shape, dtype=dC_dinps_t[ins_pos].dtype
                          )
                      )
  
              if isinstance(dC_dinps_t[ins_pos].type, NullType):
                  # We cannot use Null in the inner graph, so we
                  # use a zero tensor of the appropriate shape instead.
                  inner_out_mitmot.append(
                      aet.zeros(diff_inputs[ins_pos].shape, dtype=config.floatX)
                  )
              else:
                  inner_out_mitmot.append(dC_dinps_t[ins_pos])
  
              for _sh in self.inner_shared(self_inputs):
                  if _sh in graph_inputs([dC_dinps_t[ins_pos]]):
                      through_shared = True
  
              if isinstance(dC_dinps_t[ins_pos].type, NullType):
                  type_outs.append(dC_dinps_t[ins_pos].type.why_null)
              elif through_shared:
                  type_outs.append("through_shared")
              elif disconnected_dC_dinps_t[ins_pos]:
                  type_outs.append("disconnected")
              else:
                  type_outs.append("connected")
  
              inner_inp_mitmot += [dC_dXts[out_pos], dC_dXtm1s[ins_pos - self.n_seqs]]
              n_mitmot_outs += 1
              out_pos += 1
              ins_pos += 1
              n_mitmot_inps += 2
  
          n_nit_sot = self.n_seqs
          inner_out_nitsot = dC_dinps_t[: self.n_seqs]
          inner_out_sitsot = dC_dinps_t[ins_pos:]
          for _p, vl in enumerate(inner_out_sitsot):
              through_shared = False
              for _sh in self.inner_shared(self_inputs):
                  if _sh in graph_inputs([vl]):
                      through_shared = True
              if isinstance(vl.type, NullType):
                  type_outs.append(vl.type.why_null)
                  # Replace the inner output with a zero tensor of
                  # the right shape
                  inner_out_sitsot[_p] = aet.zeros(
                      diff_inputs[ins_pos + _p].shape, dtype=config.floatX
                  )
              elif through_shared:
                  type_outs.append("through_shared")
              elif disconnected_dC_dinps_t[_p + ins_pos]:
                  type_outs.append("disconnected")
              else:
                  type_outs.append("connected")
  
          for _p, vl in enumerate(inner_out_nitsot):
              through_shared = False
              for _sh in self.inner_shared(self_inputs):
                  if _sh in graph_inputs([vl]):
                      through_shared = True
              if isinstance(vl.type, NullType):
                  type_outs.append(vl.type.why_null)
                  # Replace the inner output with a zero tensor of
                  # the right shape
                  inner_out_nitsot[_p] = aet.zeros(
                      diff_inputs[_p].shape, dtype=config.floatX
                  )
  
              if through_shared:
                  type_outs.append("through_shared")
              elif disconnected_dC_dinps_t[_p]:
                  type_outs.append("disconnected")
              else:
                  type_outs.append("connected")
  
          inner_inp_sitsot = dC_dXtm1s[ins_pos - self.n_seqs :]
          outer_inp_sitsot = []
          for _idx, y in enumerate(inner_inp_sitsot):
              x = self.outer_non_seqs(inputs)[_idx]
              if isinstance(y.type, NullType):
                  # Cannot use dC_dXtm1s.dtype, so we use floatX instead.
                  outer_inp_sitsot.append(
                      aet.zeros(
                          [grad_steps + 1] + [x.shape[i] for i in range(x.ndim)],
                          dtype=config.floatX,
                      )
                  )
                  # replace y by a zero tensor of the right shape
                  inner_inp_sitsot[_idx] = aet.zeros(
                      diff_inputs[ins_pos + _idx].shape, dtype=config.floatX
                  )
  
              else:
                  outer_inp_sitsot.append(
                      aet.zeros(
                          [grad_steps + 1] + [x.shape[i] for i in range(x.ndim)],
                          dtype=y.dtype,
                      )
                  )
  
          n_sitsot_outs = len(outer_inp_sitsot)
          new_tap_array = mitmot_inp_taps + [[-1] for k in range(n_sitsot_outs)]
  
          info = OrderedDict()
          info["n_seqs"] = len(outer_inp_seqs)
          info["n_mit_sot"] = 0
          info["tap_array"] = new_tap_array
          info["gpua"] = False
          info["n_mit_mot"] = len(outer_inp_mitmot)
          info["n_mit_mot_outs"] = n_mitmot_outs
          info["mit_mot_out_slices"] = mitmot_out_taps
          info["truncate_gradient"] = self.truncate_gradient
          info["n_sit_sot"] = n_sitsot_outs
          info["n_shared_outs"] = 0
          info["n_nit_sot"] = n_nit_sot
          info["as_while"] = False
          info["profile"] = self.profile
          info["destroy_map"] = OrderedDict()
          if self.name:
              info["name"] = "grad_of_" + self.name
          else:
              info["name"] = None
          info["mode"] = self.mode
          info["allow_gc"] = self.allow_gc
  
          outer_inputs = (
              [grad_steps]
              + outer_inp_seqs
              + outer_inp_mitmot
              + outer_inp_sitsot
              + [n_steps if self.as_while else inputs[0] for _ in range(n_nit_sot)]
              + self.outer_shared(inputs)
              + self.outer_non_seqs(inputs)
          )
  
          inner_gfn_ins = (
              inner_inp_seqs
              + inner_inp_mitmot
              + inner_inp_sitsot
              + self.inner_shared(self_inputs)
              + self.inner_non_seqs(self_inputs)
          )
          inner_gfn_outs = inner_out_mitmot + inner_out_sitsot + inner_out_nitsot
  
          local_op = Scan(inner_gfn_ins, inner_gfn_outs, info)
          outputs = local_op(*outer_inputs)
          if type(outputs) not in (list, tuple):
              outputs = [outputs]
          # Re-order the gradients correctly
          gradients = [DisconnectedType()()]
  
          offset = self.n_mit_mot + self.n_mit_sot + self.n_sit_sot + n_sitsot_outs
          for p, (x, t) in enumerate(
              zip(
                  outputs[offset : offset + self.n_seqs],
                  type_outs[offset : offset + self.n_seqs],
              )
          ):
              if t == "connected":
                  # If the forward scan is in as_while mode, we need to pad
                  # the gradients, so that they match the size of the input
                  # sequences.
                  if self.as_while:
                      n_zeros = inputs[0] - n_steps
                      shp = (n_zeros,)
                      if x.ndim > 1:
                          shp = shp + tuple(x.shape[i] for i in range(1, x.ndim))
                      z = aet.zeros(shp, dtype=x.dtype)
                      x = aet.concatenate([x[::-1], z], axis=0)
                      gradients.append(x)
                  else:
                      gradients.append(x[::-1])
              elif t == "disconnected":
                  gradients.append(DisconnectedType()())
              elif t == "through_shared":
                  gradients.append(
                      grad_undefined(
                          self, p + 1, inputs[p + 1], "Depends on a shared variable"
                      )
                  )
              else:
                  # t contains the "why_null" string of a NullType
                  gradients.append(NullType(t)())
  
          end = self.n_mit_mot + self.n_mit_sot + self.n_sit_sot
          for p, (x, t) in enumerate(zip(outputs[:end], type_outs[:end])):
              if t == "connected":
                  # If the forward scan is in as_while mode, we need to pad
                  # the gradients, so that they match the size of the input
                  # sequences.
                  if self.as_while:
                      n_zeros = inputs[0] - grad_steps
                      shp = (n_zeros,)
                      if x.ndim > 1:
                          shp = shp + tuple(x.shape[i] for i in range(1, x.ndim))
                      z = aet.zeros(shp, dtype=x.dtype)
                      x = aet.concatenate([x[::-1], z], axis=0)
                      gradients.append(x)
                  else:
                      gradients.append(x[::-1])
              elif t == "disconnected":
                  gradients.append(DisconnectedType()())
              elif t == "through_shared":
                  gradients.append(
                      grad_undefined(
                          self,
                          p + 1 + self.n_seqs,
                          inputs[p + 1 + self.n_seqs],
                          "Depends on a shared variable",
                      )
                  )
              else:
                  # t contains the "why_null" string of a NullType
                  gradients.append(NullType(t)())
  
          start = len(gradients)
          node = outs[0].owner
          for idx in range(self.n_shared_outs):
              disconnected = True
              connected_flags = self.connection_pattern(node)[idx + start]
              for dC_dout, connected in zip(dC_douts, connected_flags):
                  if not isinstance(dC_dout.type, DisconnectedType) and connected:
                      disconnected = False
              if disconnected:
                  gradients.append(DisconnectedType()())
              else:
                  gradients.append(
                      grad_undefined(
                          self, idx, inputs[idx], "Shared Variable with update"
                      )
                  )
  
          start = len(gradients)
          gradients += [DisconnectedType()() for _ in range(self.n_nit_sot)]
          begin = end
  
          end = begin + n_sitsot_outs
          for p, (x, t) in enumerate(zip(outputs[begin:end], type_outs[begin:end])):
              if t == "connected":
                  gradients.append(x[-1])
              elif t == "disconnected":
                  gradients.append(DisconnectedType()())
              elif t == "through_shared":
                  gradients.append(
                      grad_undefined(
                          self,
                          p + begin + 1,
                          inputs[p + begin + 1],
                          "Depends on a shared variable",
                      )
                  )
              else:
                  # t contains the "why_null" string of a NullType
                  gradients.append(NullType(t)())
  
          # Mask disconnected gradients
          # Ideally we would want to assert that the gradients we are
          # replacing do indeed evaluate to 0, though that is not practical
          # from a computational point of view
          # The gradients of scan are computed replacing Disconnected with 0,
          # because through the recurrence they can become nonzero
          for idx in range(len(gradients)):
              disconnected = True
              for kdx in range(len(node.outputs)):
                  if connection_pattern[idx][kdx] and not isinstance(
                      dC_douts[kdx].type, DisconnectedType
                  ):
                      disconnected = False
              if disconnected:
                  gradients[idx] = DisconnectedType()()
          return gradients
  
      def R_op(self, inputs, eval_points):
          # Step 0. Prepare some shortcut variable
          self_inputs = self.inputs
          rop_of_inputs = (
              self_inputs[: self.n_seqs + self.n_outs]
              + self_inputs[self.n_seqs + self.n_outs + self.n_shared_outs :]
          )
          self_outputs = self.outputs
  
          # Step 1. Compute the R_op of the inner function
          inner_eval_points = [safe_new(x, "_evalpoint") for x in rop_of_inputs]
          if self.as_while:
              rop_self_outputs = self_outputs[:-1]
          else:
              rop_self_outputs = self_outputs
          if self.info["n_shared_outs"] > 0:
              rop_self_outputs = rop_self_outputs[: -self.info["n_shared_outs"]]
          rop_outs = Rop(rop_self_outputs, rop_of_inputs, inner_eval_points)
          if type(rop_outs) not in (list, tuple):
              rop_outs = [rop_outs]
          # Step 2. Figure out what corresponds to what in the scan
  
          # When doing the R-op of scan, you end up having double of each type of
          # input, because for each sequence you need also its eval point, for
          # each mit_mot, mit_sot, sit_sot or other type of inputs the same.
          # Interestingly enough, all these types of eval points behave the same
          # way as the input to which they correspond
          # The only exception is the eval point for the number of sequences, and
          # evan point for the number of nit_sot which I think should just be
          # ignored (?)
          info = OrderedDict()
          info["n_seqs"] = self.n_seqs * 2
          info["n_mit_sot"] = self.n_mit_sot * 2
          info["n_sit_sot"] = self.n_sit_sot * 2
          info["n_mit_mot"] = self.n_mit_mot * 2
          info["n_nit_sot"] = self.n_nit_sot * 2
          info["n_shared_outs"] = self.n_shared_outs
          info["gpua"] = False
          info["as_while"] = self.as_while
          info["profile"] = self.profile
          info["truncate_gradient"] = self.truncate_gradient
          if self.name:
              info["name"] = "rop_of_" + self.name
          else:
              info["name"] = None
          info["mode"] = self.mode
          info["allow_gc"] = self.allow_gc
          info["mit_mot_out_slices"] = self.mit_mot_out_slices * 2
          info["destroy_map"] = OrderedDict()
          new_tap_array = []
          b = 0
          e = self.n_mit_mot
          new_tap_array += self.tap_array[b:e] * 2
          b = e
          e += self.n_mit_sot
          new_tap_array += self.tap_array[b:e] * 2
          b = e
          e += self.n_sit_sot
          new_tap_array += self.tap_array[b:e] * 2
          info["tap_array"] = new_tap_array
  
          # Sequences ...
          b = 1
          ib = 0
          e = 1 + self.n_seqs
          ie = self.n_seqs
          clean_eval_points = []
          for inp, evp in zip(inputs[b:e], eval_points[b:e]):
              if evp is not None:
                  clean_eval_points.append(evp)
              else:
                  clean_eval_points.append(inp.zeros_like())
  
          scan_seqs = inputs[b:e] + clean_eval_points
          inner_seqs = self_inputs[ib:ie] + inner_eval_points[ib:ie]
  
          # MIT_MOT sequences ...
          b = e
          e = e + self.n_mit_mot
          ib = ie
          ie = ie + int(np.sum([len(x) for x in self.tap_array[: self.n_mit_mot]]))
          clean_eval_points = []
          for inp, evp in zip(inputs[b:e], eval_points[b:e]):
              if evp is not None:
                  clean_eval_points.append(evp)
              else:
                  clean_eval_points.append(inp.zeros_like())
  
          scan_mit_mot = inputs[b:e] + clean_eval_points
          inner_mit_mot = self_inputs[ib:ie] + inner_eval_points[ib:ie]
  
          # MIT_SOT sequences ...
          b = e
          e = e + self.n_mit_sot
          ib = ie
          ie = ie + int(
              np.sum(
                  [
                      len(x)
                      for x in self.tap_array[
                          self.n_mit_mot : self.n_mit_mot + self.n_mit_sot
                      ]
                  ]
              )
          )
          clean_eval_points = []
          for inp, evp in zip(inputs[b:e], eval_points[b:e]):
              if evp is not None:
                  clean_eval_points.append(evp)
              else:
                  clean_eval_points.append(inp.zeros_like())
  
          scan_mit_sot = inputs[b:e] + eval_points[b:e]
          inner_mit_sot = self_inputs[ib:ie] + inner_eval_points[ib:ie]
  
          # SIT_SOT sequences ...
          b = e
          e = e + self.n_sit_sot
          ib = ie
          ie = ie + self.n_sit_sot
          clean_eval_points = []
          for inp, evp in zip(inputs[b:e], eval_points[b:e]):
              if evp is not None:
                  clean_eval_points.append(evp)
              else:
                  clean_eval_points.append(inp.zeros_like())
  
          scan_sit_sot = inputs[b:e] + clean_eval_points
          inner_sit_sot = self_inputs[ib:ie] + inner_eval_points[ib:ie]
  
          # Shared outs ...
          b = e
          e = e + self.n_shared_outs
          ib = ie
          ie = ie + self.n_shared_outs
          scan_shared = inputs[b:e]
          inner_shared = self_inputs[ib:ie]
  
          # NIT_SOT sequences
          b = e
          e = e + self.n_nit_sot
          scan_nit_sot = inputs[b:e] * 2
  
          # All other arguments
          clean_eval_points = []
          for inp, evp in zip(inputs[e:], eval_points[e:]):
              if evp is not None:
                  clean_eval_points.append(evp)
              else:
                  clean_eval_points.append(inp.zeros_like())
          scan_other = inputs[e:] + clean_eval_points
          # inner_eval_points do not have entries for shared variables
          inner_other = self_inputs[ie:] + inner_eval_points[ib:]
  
          # Outputs
          n_mit_mot_outs = int(np.sum([len(x) for x in self.mit_mot_out_slices]))
          info["n_mit_mot_outs"] = n_mit_mot_outs * 2
          b = 0
          e = n_mit_mot_outs
          inner_out_mit_mot = self_outputs[b:e] + rop_outs[b:e]
          b = e
          e = e + self.n_mit_sot
          inner_out_mit_sot = self_outputs[b:e] + rop_outs[b:e]
          b = e
          e = e + self.n_sit_sot
          inner_out_sit_sot = self_outputs[b:e] + rop_outs[b:e]
          b = e
          e = e + self.n_nit_sot
          inner_out_nit_sot = self_outputs[b:e] + rop_outs[b:e]
          b = e
          e = e + self.n_shared_outs
          inner_out_shared = self_outputs[b:e]
  
          inner_ins = (
              inner_seqs
              + inner_mit_mot
              + inner_mit_sot
              + inner_sit_sot
              + inner_shared
              + inner_other
          )
          inner_outs = (
              inner_out_mit_mot
              + inner_out_mit_sot
              + inner_out_sit_sot
              + inner_out_nit_sot
              + inner_out_shared
          )
  
          if self.as_while:
              inner_outs += [self_outputs[-1]]
          scan_inputs = (
              [inputs[0]]
              + scan_seqs
              + scan_mit_mot
              + scan_mit_sot
              + scan_sit_sot
              + scan_shared
              + scan_nit_sot
              + scan_other
          )
  
          local_op = Scan(inner_ins, inner_outs, info)
          outputs = local_op(*scan_inputs)
          if type(outputs) not in (list, tuple):
              outputs = [outputs]
          # Select only the result of the R_op results
          final_outs = []
          b = self.n_mit_mot
          e = self.n_mit_mot * 2
          final_outs += outputs[b:e]
          b = e + self.n_mit_sot
          e = e + self.n_mit_sot * 2
          final_outs += outputs[b:e]
          b = e + self.n_sit_sot
          e = e + self.n_sit_sot * 2
          final_outs += outputs[b:e]
          b = e + self.n_nit_sot
          e = e + self.n_nit_sot * 2
          final_outs += outputs[b:e]
          final_outs += [None] * self.n_shared_outs
  
          return final_outs
  
  
  # Since Scan is an op that contains an Aesara compiled function, it is
  # useful to let DebugMode know about it.
  ops_with_inner_function[Scan] = "fn"
  
  
  @register_profiler_printer
  def profile_printer(
      message, compile_time, fct_call_time, apply_time, apply_cimpl, outputs_size, file
  ):
      # Scan overhead profile
      if any(
          [
              isinstance(node.op, Scan) and v > 0
              for (fgraph, node), v in apply_time.items()
          ]
      ):
          print("", file=file)
          print("Scan overhead:", file=file)
          print(
              "<Scan op time(s)> <sub scan fct time(s)> <sub scan op "
              "time(s)> <sub scan fct time(% scan op time)> <sub scan "
              "op time(% scan op time)> <node>",
              file=file,
          )
  
          total_super_scan_time = 0
          total_scan_fct_time = 0
          total_scan_op_time = 0
          for (fgraph, node), v in apply_time.items():
              if isinstance(node.op, Scan) and not node.op.fn.profile:
                  print(
                      "  One scan node do not have its inner profile enabled. "
                      "If you enable Aesara profiler with "
                      "'aesara.function(..., profile=True)', you must manually"
                      " enable the profiling for each scan too: "
                      "'aesara.scan(...,profile=True)'."
                      " Or use Aesara flag 'profile=True'.",
                      file=file,
                  )
              elif isinstance(node.op, Scan) and node.op.fn.profile:
                  if v > 0:
                      scan_fct_time = node.op.fn.profile.call_time
                      scan_op_time = sum(node.op.fn.profile.apply_time.values())
                      total_super_scan_time += v
                      total_scan_fct_time += scan_fct_time
                      total_scan_op_time += scan_op_time
                      print(
                          "      %5.1fs  %5.1fs  %5.1fs  %5.1f%%  %5.1f%%"
                          % (
                              v,
                              scan_fct_time,
                              scan_op_time,
                              scan_fct_time / v * 100,
                              scan_op_time / v * 100,
                          ),
                          node,
                          file=file,
                      )
                  else:
                      print(
                          (" The node took 0s, so we can not " "compute the overhead"),
                          node,
                          file=file,
                      )
          if total_super_scan_time == 0:
              print("  No scan have its inner profile enabled.", file=file)
          else:
              print(
                  "total %5.1fs  %5.1fs  %5.1fs  %5.1f%%  %5.1f%%"
                  % (
                      total_super_scan_time,
                      total_scan_fct_time,
                      total_scan_op_time,
                      total_scan_fct_time / total_super_scan_time * 100,
                      total_scan_op_time / total_super_scan_time * 100,
                  ),
                  file=file,
              )
   import builtins
  import warnings
  
  import numpy as np
  
  from aesara import config, printing
  from aesara import scalar as aes
  from aesara.gradient import DisconnectedType
  from aesara.graph.basic import Apply, Variable
  from aesara.graph.op import COp, Op
  from aesara.graph.params_type import ParamsType
  from aesara.graph.type import Generic
  from aesara.misc.safe_asarray import _asarray
  from aesara.printing import pprint
  from aesara.scalar.basic import BinaryScalarOp
  from aesara.tensor.basic import (
      alloc,
      arange,
      as_tensor_variable,
      cast,
      concatenate,
      constant,
      patternbroadcast,
      stack,
      switch,
  )
  from aesara.tensor.elemwise import (
      CAReduce,
      CAReduceDtype,
      DimShuffle,
      Elemwise,
      scalar_elemwise,
  )
  from aesara.tensor.shape import shape
  from aesara.tensor.type import (
      complex_dtypes,
      continuous_dtypes,
      discrete_dtypes,
      int_dtypes,
      integer_dtypes,
      tensor,
      uint_dtypes,
  )
  from aesara.tensor.type_other import NoneConst
  from aesara.tensor.utils import as_list
  from aesara.tensor.var import TensorConstant, _tensor_py_operators
  
  
  # We capture the builtins that we are going to replace to follow the numpy API
  _abs = builtins.abs
  
  
  if int(config.tensor__cmp_sloppy) > 1:
      # This config variable is a quick-and-dirty way to get low-precision
      # comparisons.  For a more precise setting of these tolerances set
      # them explicitly in your user code by assigning, for example,
      # "aesara.tensor.math.float32_atol = ..."
  
      # When config.tensor__cmp_sloppy>1 we are even more sloppy. This is
      # useful to test the GPU as they don't use extended precision and
      # this cause some difference bigger then the normal sloppy.
      float16_atol = 1e-2
      float16_rtol = 5e-2
  
      float32_atol = 5e-4
      float32_rtol = 1e-3
  
      float64_rtol = 1e-4
      float64_atol = 1e-3
  elif int(config.tensor__cmp_sloppy):
      float16_atol = 5e-3
      float16_rtol = 1e-2
  
      float32_atol = 1e-4
      float32_rtol = 1e-3
  
      float64_rtol = 1e-4
      float64_atol = 1e-3
  else:
      # If you change those value in test don't forget to put them back
      # when the test end.  Don't forget the case when the test fail.
      float16_atol = 1e-3
      float16_rtol = 1e-3
  
      float32_atol = 1e-5
      float32_rtol = 1e-5
  
      # defaults in numpy.allclose
      # Don't be more strict then numpy rtol
      # It cause useless error.
      float64_rtol = 1.0000000000000001e-05
      float64_atol = 1e-8
  
  
  def _get_atol_rtol(a, b):
      tiny = ("float16",)
      narrow = ("float32", "complex64")
      if (str(a.dtype) in tiny) or (str(b.dtype) in tiny):
          atol = float16_atol
          rtol = float16_rtol
      elif (str(a.dtype) in narrow) or (str(b.dtype) in narrow):
          atol = float32_atol
          rtol = float32_rtol
      else:
          atol = float64_atol
          rtol = float64_rtol
      return atol, rtol
  
  
  def _allclose(a, b, rtol=None, atol=None):
      a = np.asarray(a)
      b = np.asarray(b)
      atol_, rtol_ = _get_atol_rtol(a, b)
      if rtol is not None:
          rtol_ = rtol
      if atol is not None:
          atol_ = atol
  
      return np.allclose(a, b, atol=atol_, rtol=rtol_)
  
  
  class MaxAndArgmax(COp):
      """
      Calculate the max and argmax over a given axis or over all axes.
  
      """
  
      nin = 2  # tensor, axis
      nout = 2  # max val, max idx
      E_axis = "invalid axis"
      params_type = Generic()
      __props__ = ("axis",)
      _f16_ok = True
  
      def __init__(self, axis):
          assert isinstance(axis, list)
          self.axis = tuple(axis)
  
      def get_params(self, node):
          return self.axis
  
      def make_node(self, x):
          x = as_tensor_variable(x)
  
          # We keep the original broadcastable flags for dimensions on which
          # we do not perform the max / argmax.
          all_axes = set(self.axis)
          broadcastable = [
              b for i, b in enumerate(x.type.broadcastable) if i not in all_axes
          ]
          inputs = [x]
          outputs = [
              tensor(x.type.dtype, broadcastable, name="max"),
              tensor("int64", broadcastable, name="argmax"),
          ]
          return Apply(self, inputs, outputs)
  
      def perform(self, node, inp, outs, params):
          x = inp[0]
          axes = params
          max, max_idx = outs
          if axes is None:
              axes = tuple(range(x.ndim))
          else:
              axes = tuple(int(ax) for ax in axes)
          max[0] = _asarray(np.max(x, axes), dtype=node.outputs[0].dtype)
          # Numpy does not support multiple axes for argmax
          # Work around
          keep_axes = np.array([i for i in range(x.ndim) if i not in axes], dtype="int64")
          # Not-reduced axes in front
          transposed_x = np.transpose(x, np.concatenate((keep_axes, axes)))
          kept_shape = transposed_x.shape[: len(keep_axes)]
          reduced_shape = transposed_x.shape[len(keep_axes) :]
  
          # Numpy.prod returns 1.0 when arg is empty, so we cast it to int64
          # Otherwise reshape would complain citing float arg
          new_shape = kept_shape + (np.prod(reduced_shape, dtype="int64"),)
          reshaped_x = transposed_x.reshape(new_shape)
  
          max_idx[0] = _asarray(np.argmax(reshaped_x, axis=-1), dtype="int64")
  
      def c_code(self, node, name, inp, out, sub):
          if len(self.axis) != 1 and len(self.axis) != node.inputs[0].ndim:
              raise NotImplementedError(
                  "NumPy C-API can compute max and argmax only for 1 axis or for all axes."
              )
          x = inp[0]
          axis = sub["params"]
          max, argmax = out
          fail = sub["fail"]
          ret = """
          #if PY_MAJOR_VERSION >= 3
              #ifndef PyInt_AS_LONG
                  #define PyInt_AS_LONG PyLong_AS_LONG
              #endif
          #endif
  
          int axis;
  
          if (PyTuple_GET_SIZE(%(axis)s) == PyArray_NDIM(%(x)s)) {
              axis = NPY_MAXDIMS;
          } else if(PyTuple_GET_SIZE(%(axis)s) == 1) {
              PyObject* axis_object = PyTuple_GET_ITEM(%(axis)s, 0);
              axis = (int)PyInt_AS_LONG(axis_object);
              if (axis > PyArray_NDIM(%(x)s)-1 || axis < -PyArray_NDIM(%(x)s)) {
                  PyErr_SetString(PyExc_ValueError,
                  "MaxAndArgmax: bad axis argument");
                  %(fail)s
              }
          } else {
              PyErr_SetString(PyExc_NotImplementedError,
              "MaxAndArgmax: NumPy C-API can compute max and argmax only for 1 axis or for all axes.");
              %(fail)s
          }
  
          Py_CLEAR(%(max)s);
          Py_CLEAR(%(argmax)s);//todo pass them as out parameter.
  
          %(max)s = (PyArrayObject*)PyArray_Max(%(x)s, axis, NULL);
          if (%(max)s == NULL) {
              %(fail)s;
          }
          if (!PyArray_CheckExact(%(max)s)) {
              %(max)s = (PyArrayObject*)PyArray_FromAny((PyObject*)%(max)s, NULL, 0, 0, NPY_ARRAY_ENSUREARRAY, NULL);
              if(%(max)s == NULL){
                  %(fail)s;
              }
          }
  
          %(argmax)s = (PyArrayObject*)PyArray_ArgMax(%(x)s, axis, NULL);
          if (%(argmax)s == NULL) {
              Py_CLEAR(%(max)s);
              %(fail)s;
          }
          if (!PyArray_CheckExact(%(argmax)s)) {
              %(argmax)s = (PyArrayObject*)PyArray_FromAny((PyObject*)%(argmax)s, NULL, 0, 0, NPY_ARRAY_ENSUREARRAY, NULL);
              if(%(argmax)s == NULL){
                  %(fail)s;
              }
          }
          if (PyArray_TYPE(%(argmax)s) != NPY_INT64) {
              PyObject * tmp = PyArray_Cast(%(argmax)s, NPY_INT64);
              if (NULL == tmp){
                  %(fail)s;
              }
              Py_DECREF(%(argmax)s);
              %(argmax)s = (PyArrayObject*)tmp;
          }
          """
          return ret % locals()
  
      def c_code_cache_version(self):
          return (5,)
  
      def infer_shape(self, fgraph, node, shapes):
          ishape = shapes[0]
          rval = tuple(
              ishape[i]
              for (i, b) in enumerate(node.inputs[0].type.broadcastable)
              if i not in self.axis
          )
          return [rval, rval]
  
      def R_op(self, inputs, eval_points):
          if eval_points[0] is None:
              return [None, None]
          if len(self.axis) != 1:
              raise ValueError("R_op supported for arg_max only for " "one axis!")
          if self.axis[0] > 1:
              raise ValueError("R_op supported for arg_max only when " " axis is 0 or 1")
          if inputs[0].ndim != 2:
              raise ValueError(
                  "R_op supported for arg_max only when " " input is a matrix"
              )
          max_vals, max_pos = self.make_node(*inputs).outputs
          if self.axis[0] == 0:
              return [eval_points[0][max_pos, arange(eval_points[0].shape[1])], None]
          else:
              return [eval_points[0][arange(eval_points[0].shape[0]), max_pos], None]
  
      def grad(self, inp, grads):
          # The strict sense mathematical gradient of the maximum function is
          # not calculated here for it is not defined at every point where some
          # coordinates are identical. However, since the latter set has null
          # Lebesgue measure, the result may be interpreted as weak gradient.
  
          # @note: This function should work correctly for L{vector}s.
          # (x, y), (gz, gw)
          # gz*dz/dx + gw*dw/dx, gz*dz/dy + gw*dw/dy
          # gMax * dMax/dx + gArgMax * dArgMax/dx,
          # gMax * dMax/daxis + gArgMax * dArgMax/daxis
          # g_max has one less dimension than x, so you need to complete
          # g_max to x's shape when axis=0 the broadcasting mechanism
          # does it automatically
          x = inp[0]
          axis = as_tensor_variable(self.axis)
          g_max, g_max_idx = grads
  
          g_max_disconnected = isinstance(g_max.type, DisconnectedType)
          g_max_idx_disconnected = isinstance(g_max_idx.type, DisconnectedType)
  
          # if the op is totally disconnected, so are its inputs
          if g_max_disconnected and g_max_idx_disconnected:
              return [DisconnectedType()(), DisconnectedType()()]
  
          # if the max is disconnected but the argmax is not,
          # the gradient on its inputs is zero
          if g_max_disconnected:
              return [x.zeros_like()]
          if NoneConst.equals(axis):
              axis_ = list(range(x.ndim))
          else:
              axis_ = axis
          xmax = max(x, axis_)
  
          # Raise the g_max and xmax to the same number of dim as the input.
          pattern = []
          out_dim = 0
          if NoneConst.equals(axis):
              # We are taking the max/argmax over all dimensions.
              axis = None
          for i in range(x.ndim):
              if axis is None or i in axis.data:
                  pattern.append("x")
              else:
                  pattern.append(out_dim)
                  out_dim += 1
          g_max_pad = DimShuffle(g_max.broadcastable, pattern)(g_max)
          xmax_pad = DimShuffle(xmax.broadcastable, pattern)(xmax)
  
          # Set the grad to the correct position.
          g_x = eq(xmax_pad, x) * g_max_pad
          return (g_x,)
  
  
  class Argmax(COp):
      """
      Calculate the argmax over a given axis or over all axes.
      """
  
      nin = 2  # tensor, axis
      nout = 1
      E_axis = "invalid axis"
      __props__ = ("axis",)
      _f16_ok = True
  
      params_type = ParamsType(c_axis=aes.int64)
  
      def __init__(self, axis):
          if axis is not None:
              axis = tuple(axis)
          self.axis = tuple(axis)
  
      def get_params(self, node):
          if self.axis is not None and len(self.axis) == 1:
              c_axis = np.int64(self.axis[0])
          else:
              # The value here doesn't matter, it won't be used
              c_axis = np.int64(-1)
          return self.params_type.get_params(c_axis=c_axis)
  
      def make_node(self, x, axis=None):
          x = as_tensor_variable(x)
          if self.axis is None:
              all_axes = list(range(x.ndim))
          else:
              all_axes = self.axis
          inputs = [x]
  
          # We keep the original broadcastable flags for dimensions on which
          # we do not perform the argmax.
          broadcastable = [
              b for i, b in enumerate(x.type.broadcastable) if i not in all_axes
          ]
          outputs = [tensor("int64", broadcastable, name="argmax")]
          return Apply(self, inputs, outputs)
  
      def prepare_node(self, node, storage_map, compute_map, impl):
          if len(node.inputs) == 2:
              raise ValueError(
                  "You are trying to compile a graph with an old Argmax node.  Either reoptimize your graph or rebuild it to get the new node format."
              )
  
      def perform(self, node, inp, outs, params):
          (x,) = inp
          axes = self.axis
          (max_idx,) = outs
          if axes is None:
              axes = tuple(range(x.ndim))
  
          # Numpy does not support multiple axes for argmax
          # Work around
          keep_axes = np.array([i for i in range(x.ndim) if i not in axes], dtype="int64")
          # Not-reduced axes in front
          transposed_x = np.transpose(x, np.concatenate((keep_axes, axes)))
          kept_shape = transposed_x.shape[: len(keep_axes)]
          reduced_shape = transposed_x.shape[len(keep_axes) :]
          new_shape = kept_shape + (np.prod(reduced_shape),)
          reshaped_x = transposed_x.reshape(new_shape)
  
          max_idx[0] = _asarray(np.argmax(reshaped_x, axis=-1), dtype="int64")
  
      def c_code(self, node, name, inp, out, sub):
          (x,) = inp
          (argmax,) = out
          fail = sub["fail"]
          params = sub["params"]
          if self.axis is None:
              axis_code = "axis = NPY_MAXDIMS;"
          else:
              if len(self.axis) > 1:
                  raise NotImplementedError()
              # params is only used here for now
              axis_code = (
                  """
              axis = %(params)s->c_axis;
              if(axis > PyArray_NDIM(%(x)s)-1 || axis < -PyArray_NDIM(%(x)s)){
                  PyErr_SetString(PyExc_ValueError,
                  "Argmax, bad axis argument");
                  %(fail)s
              }
              """
                  % locals()
              )
          ret = """
          int axis;
  
          Py_CLEAR(%(argmax)s);//todo pass them as out parameter.
          %(axis_code)s
  
          %(argmax)s = (PyArrayObject*)PyArray_ArgMax(%(x)s, axis, NULL);
          if(%(argmax)s == NULL){
              %(fail)s;
          }
          if(!PyArray_CheckExact(%(argmax)s)){
              %(argmax)s = (PyArrayObject*)PyArray_FromAny((PyObject*)%(argmax)s, NULL, 0, 0, NPY_ARRAY_ENSUREARRAY, NULL);
              if(%(argmax)s == NULL){
                  %(fail)s;
              }
          }
          if(PyArray_TYPE(%(argmax)s) != NPY_INT64){
              PyObject * tmp = PyArray_Cast(%(argmax)s, NPY_INT64);
              if (NULL == tmp){
                  %(fail)s;
              }
              Py_DECREF(%(argmax)s);
              %(argmax)s = (PyArrayObject*)tmp;
          }
          """
          return ret % locals()
  
      def c_code_cache_version(self):
          return (1,)
  
      def infer_shape(self, fgraph, node, shapes):
          (ishape,) = shapes
          if self.axis is None:
              return [()]
          rval = tuple(
              [
                  ishape[i]
                  for (i, b) in enumerate(node.inputs[0].type.broadcastable)
                  if i not in self.axis
              ]
          )
          return [rval]
  
      def grad(self, inp, grads):
          (x,) = inp
  
          return [x.zeros_like()]
  
  
  def makeKeepDims(x, y, axis):
      """
      Reintroduces in y with length one the axes of x which have been left out
      in a prior reduction of x. With this option, the resulting tensor will
      broadcast correctly against the original tensor x.
  
      """
      x = as_tensor_variable(x)
      y = as_tensor_variable(y)
  
      if axis is None:
          axis = list(range(x.type.ndim))
      elif isinstance(axis, (int, np.integer)):
          axis = [axis]
      elif isinstance(axis, np.ndarray) and axis.ndim == 0:
          axis = [int(axis)]
      else:
          axis = [int(a) for a in axis]
      newaxis = []
      for a in axis:
          if not isinstance(a, int):
              raise ValueError("keepdims option can be used only with constant axis")
          if a < 0:
              a += x.type.ndim
          newaxis.append(a)
      i = 0
      new_dims = []
      for j, _ in enumerate(x.type.broadcastable):
          if j in newaxis:
              new_dims.append("x")
          else:
              new_dims.append(i)
              i += 1
      return DimShuffle(y.type.broadcastable, new_dims)(y)
  
  
  def check_and_normalize_axes(x, axis):
      """Check axes, normalize and convert them to a Python list of integers.
  
      Parameters
      ----------
      x: TensorVariable
      axis: int, tuple or list of integers
  
      Returns
      -------
      axis: list of integers
          Return an empty list if argument is None.
  
      """
      x = as_tensor_variable(x)
      if axis is None:
          axis = []
      elif isinstance(axis, (int, np.integer)) or (
          isinstance(axis, np.ndarray) and axis.ndim == 0
      ):
          axis = [int(axis)]
      elif isinstance(axis, (tuple, list, np.ndarray)):
          axis = [int(i) for i in axis]
      elif isinstance(axis, Variable):
          if NoneConst.equals(axis):
              axis = []
          elif not isinstance(axis, TensorConstant):
              raise TypeError(f"Computation needs a constant axis. Got {axis}")
          else:
              assert axis.dtype in integer_dtypes
              if isinstance(axis.data, (int, np.integer)) or (
                  isinstance(axis.data, np.ndarray) and axis.data.ndim == 0
              ):
                  axis = [int(axis.data)]
              elif isinstance(axis.data, (list, np.ndarray)):
                  axis = [int(i) for i in axis.data]
      else:
          raise TypeError(
              f"Axis must be an integer, tuple, list of integers or a TensorVariable. Got {axis}"
          )
      if len(axis) > 0:
          for i in range(len(axis)):
              if axis[i] < 0:
                  axis[i] += x.type.ndim
              if axis[i] < 0 or axis[i] >= x.type.ndim:
                  raise ValueError(
                      f"Computation needs a valid axis number for {int(x.type.ndim)}-D tensor. Got {int(axis[i])}"
                  )
          axis = list(set(axis))
          axis.sort()
      return axis
  
  
  def max_and_argmax(a, axis=None, keepdims=False):
      """
      Returns maximum elements and their indices obtained by iterating over
      given axis.
  
      When axis is None (the default value), the max is performed
      over the flattened tensor.
  
      Parameters
      ----------
      keepdims : bool
          If this is set to True, the axes which are reduced are left in
          the result as dimensions with size one. With this option, the result
          will broadcast correctly against the original tensor.
  
      """
      # Check axis and convert it to a Python list of integers.
      # Axis will be used as an op param of MaxAndArgmax.
      a = as_tensor_variable(a)
      axis = check_and_normalize_axes(a, axis)
      if len(axis) == 0:
          axis = list(range(a.type.ndim))
      out, argout = MaxAndArgmax(axis)(a)
  
      if keepdims:
          out = makeKeepDims(a, out, axis)
          argout = makeKeepDims(a, argout, axis)
      return [out, argout]
  
  
  class NonZeroCAReduce(CAReduce):
      def _c_all(self, node, name, inames, onames, sub):
          decl, checks, alloc, loop, end = super()._c_all(node, name, inames, onames, sub)
  
          # We add an additional check for zero-sized dimensions (This seems like
          # something that could enabled in `elemwise_cgen.make_checks`.)
          iname = inames[0]
  
          axis = self.axis
          if axis is None:
              axis = list(range(len(node.inputs[0].type.broadcastable)))
  
          pattern = [0] * len(node.inputs[0].broadcastable)
          for i in axis:
              pattern[i] = 1
  
          pattern_ = str(pattern)[1:-1]
  
          decl += f"""int tosum[]={{{pattern_}}};"""
          alloc += f"""
                  for(int i=0;i<PyArray_NDIM({iname});i++){{
                      if(PyArray_DIMS({iname})[i]==0 && tosum[i]){{
                          PyErr_Format(PyExc_ValueError,
                              "Input of CAReduce{{{node.op.scalar_op}}} has zero-size on axis %%d",i);
                          {sub["fail"]};
                      }}
                  }}
                  """
          return decl, checks, alloc, loop, end
  
  
  class Max(NonZeroCAReduce):
      nfunc_spec = ("max", 1, 1)
  
      def __init__(self, axis):
          super().__init__(aes.scalar_maximum, axis)
  
  
  class Min(NonZeroCAReduce):
      nfunc_spec = ("min", 1, 1)
  
      def __init__(self, axis):
          super().__init__(aes.scalar_minimum, axis)
  
  
  def max(x, axis=None, keepdims=False):
      """
      Returns maximum elements obtained by iterating over given axis.
  
      When axis is None (the default value), the max is performed
      over the flattened tensor.
  
      Parameters
      ----------
      keepdims: bool
          If this is set to True, the axes which are reduced are left in
          the result as dimensions with size one. With this option, the result
          will broadcast correctly against the original tensor.
  
      Notes
      -----
      We return an error as numpy when we reduce a dim with a shape of 0.
  
      """
  
      # We have a choice of implementing this call with the
      # CAReduce op or the MaxAndArgmax op.
  
      # MaxAndArgmax supports grad and Rop, so we prefer to use that.
      # CAReduce is faster, but optimizations will replace MaxAndArgmax[0]
      # with CAReduce at compile time, so at this stage the important
      # thing is supporting all user interface features, not speed.
      # Some cases can be implemented only with CAReduce.
  
      # We thus prefer to use MaxAndArgmax, if possible. It does not
      # support all axis arguments, so we may need to fall back to CAReduce.
  
      try:
          out = max_and_argmax(x, axis)[0]
      except Exception:
          out = Max(axis)(x)
  
      if keepdims:
          out = makeKeepDims(x, out, axis)
      return out
  
  
  def argmax(x, axis=None, keepdims=False):
      """
      Returns indices of maximum elements obtained by iterating over given axis.
  
      When axis is None (the default value), the argmax is performed
      over the flattened tensor.
  
      Parameters
      ----------
      keepdims : bool
          If this is set to True, the axes which are reduced are left in
          the result as dimensions with size one. With this option, the result
          will broadcast correctly against the original tensor.
  
      """
      argout = max_and_argmax(x, axis)[1]
  
      if keepdims:
          argout = makeKeepDims(x, argout, axis)
      return argout
  
  
  def min(x, axis=None, keepdims=False):
      """
      Returns minimum elements obtained by iterating over given axis.
  
      When axis is None (the default value), the min is performed
      over the flattened tensor.
  
      Parameters
      ----------
      keepdims: bool
          If this is set to True, the axes which are reduced are left in
          the result as dimensions with size one. With this option, the result
          will broadcast correctly against the original tensor.
  
      """
      x = as_tensor_variable(x)
      str_x_type = str(x.dtype)
      if str_x_type.startswith("float") or str_x_type in int_dtypes:
          return -max(-x, axis=axis, keepdims=keepdims)
      elif str_x_type in uint_dtypes:
          itype = np.iinfo(x.dtype)
          max_val = np.array(itype.max, dtype=itype.dtype)
          return max_val - max(max_val - x, axis=axis, keepdims=keepdims)
      elif str_x_type == "bool":
          return ~max(~x, axis=axis, keepdims=keepdims)
      else:
          # Be careful about unsigned integers, complex
          raise NotImplementedError()
  
  
  def argmin(x, axis=None, keepdims=False):
      """
      Returns indices of minimum elements obtained by iterating over given axis.
  
      When axis is None (the default value), the argmin is performed
      over the flattened tensor.
  
      Parameters
      ----------
      keepdims: bool
          If this is set to True, the axes which are reduced are left in
          the result as dimensions with size one. With this option, the result
          will broadcast correctly against the original tensor.
  
      """
      x = as_tensor_variable(x)
      str_x_type = str(x.dtype)
      if str_x_type.startswith("float") or str_x_type in int_dtypes:
          return argmax(-x, axis=axis, keepdims=keepdims)
      elif str_x_type in uint_dtypes:
          itype = np.iinfo(x.dtype)
          return argmax(itype.max - x, axis=axis, keepdims=keepdims)
      elif str_x_type == "bool":
          return argmax(~x, axis=axis, keepdims=keepdims)
      else:
          # Be careful about unsigned integers, complex
          raise NotImplementedError()
  
  
  def smallest(*args):
      """
      Return the [elementwise] smallest of a variable number of arguments.
  
      Like python's min.
  
      """
      if len(args) == 2:
          a, b = args
          return switch(a < b, a, b)
      else:
          return min(stack(args), axis=0)
  
  
  def largest(*args):
      """
      Return the [elementwise] largest of a variable number of arguments.
  
      Like python's max.
  
      """
      if len(args) == 2:
          a, b = args
          return switch(a > b, a, b)
      else:
          return max(stack(args), axis=0)
  
  
  @scalar_elemwise
  def lt(a, b):
      """a < b"""
  
  
  @scalar_elemwise
  def gt(a, b):
      """a > b"""
  
  
  @scalar_elemwise
  def le(a, b):
      """a <= b"""
  
  
  @scalar_elemwise
  def ge(a, b):
      """a >= b"""
  
  
  @scalar_elemwise
  def eq(a, b):
      """a == b"""
  
  
  @scalar_elemwise
  def neq(a, b):
      """a != b"""
  
  
  @scalar_elemwise
  def isnan(a):
      """isnan(a)"""
  
  
  # Rename isnan to isnan_ to allow to bypass it when not needed.
  # glibc 2.23 don't allow isnan on int, so we remove it from the graph.
  isnan_ = isnan
  
  
  def isnan(a):
      """isnan(a)"""
      a = as_tensor_variable(a)
      if a.dtype in discrete_dtypes:
          return alloc(
              np.asarray(False, dtype="bool"), *[a.shape[i] for i in range(a.ndim)]
          )
      return isnan_(a)
  
  
  @scalar_elemwise
  def isinf(a):
      """isinf(a)"""
  
  
  # Rename isnan to isnan_ to allow to bypass it when not needed.
  # glibc 2.23 don't allow isnan on int, so we remove it from the graph.
  isinf_ = isinf
  
  
  def isinf(a):
      """isinf(a)"""
      a = as_tensor_variable(a)
      if a.dtype in discrete_dtypes:
          return alloc(
              np.asarray(False, dtype="bool"), *[a.shape[i] for i in range(a.ndim)]
          )
      return isinf_(a)
  
  
  def allclose(a, b, rtol=1.0e-5, atol=1.0e-8, equal_nan=False):
      """
      Implement Numpy's ``allclose`` on tensors.
  
      ``absolute(a - b) <= (atol + rtol * absolute(b))``
  
      Parameters
      ----------
      a : tensor
          Input to compare.
      b : tensor
          Input to compare.
      rtol : float
          The relative tolerance parameter.
      atol : float
          The absolute tolerance parameter.
      equal_nan: bool
          Whether to consider nan's in the same place to be close.
  
      Returns
      -------
      bool
          A boolean value (of type int8 returned by the tensor elementwise `all`
          function) whether all elements in a and b are in the tolerance range
          defined above.
  
      Notes
      -----
      Not a symmetric equation. See Numpy's documentation.
  
      """
      return all(isclose(a, b, rtol, atol, equal_nan))
  
  
  def isclose(a, b, rtol=1.0e-5, atol=1.0e-8, equal_nan=False):
      """
      Implements Numpy's ``isclose`` on tensors.
  
      The tolerance values are positive, typically very small numbers. The
      relative difference (`rtol` * abs(`b`)) and the absolute difference
      `atol` are added together to compare against the absolute difference
      between `a` and `b`.
  
      ``absolute(a - b) <= (atol + rtol * absolute(b))``
  
      Parameters
      ----------
      a : tensor
          Input to compare.
      b : tensor
          Input to compare.
      rtol : float
          The relative tolerance parameter.
      atol : float
          The absolute tolerance parameter.
      equal_nan : bool
          Whether to consider nan's in the same place to be close
  
      Returns
      -------
      int8
          A boolean (int8) array where two arrays are element-wise equal
          within a tolerance.
  
      Notes
      -----
      Not a symmetric equation. See Numpy's documentation.
  
      Examples
      --------
      >>> import aesara
      >>> import numpy as np
      >>> a = _asarray([1e10, 1e-7], dtype="float64")
      >>> b = _asarray([1.00001e10, 1e-8], dtype="float64")
      >>> aesara.tensor.isclose(a, b).eval()
      array([1, 0], dtype=int8)
      >>> a = _asarray([1e10, 1e-8], dtype="float64")
      >>> b = _asarray([1.00001e10, 1e-9], dtype="float64")
      >>> aesara.tensor.isclose(a, b).eval()
      array([1, 1], dtype=int8)
      >>> a = _asarray([1e10, 1e-8], dtype="float64")
      >>> b = _asarray([1.0001e10, 1e-9], dtype="float64")
      >>> aesara.tensor.isclose(a, b).eval()
      array([0, 1], dtype=int8)
      >>> a = _asarray([1.0, np.nan], dtype="float64")
      >>> b = _asarray([1.0, np.nan], dtype="float64")
      >>> aesara.tensor.isclose(a, b).eval()
      array([1, 0], dtype==int8)
      >>> a = _asarray([1.0, np.nan], dtype="float64")
      >>> b = _asarray([1.0, np.nan], dtype="float64")
      >>> aesara.tensor.isclose(a, b, equal_nan=True).eval()
      array([1, 1], dtype==int8)
      >>> a = _asarray([1.0, np.inf], dtype="float64")
      >>> b = _asarray([1.0, -np.inf], dtype="float64")
      >>> aesara.tensor.isclose(a, b).eval()
      array([1, 0], dtype==int8)
      >>> a = _asarray([1.0, np.inf], dtype="float64")
      >>> b = _asarray([1.0, np.inf], dtype="float64")
      >>> aesara.tensor.isclose(a, b).eval()
      array([1, 1], dtype==int8)
  
      """
      # close will be an int8 array of 1 where within tolerance
      # and 0 where not within tolerance or there was a nan or inf value.
      diff = _abs(a - b)
      tolerance = atol + rtol * _abs(b)
      close_prelim = le(diff, tolerance)
  
      a_nan = isnan(a)
      b_nan = isnan(b)
      nans = bitwise_or(a_nan, b_nan)
  
      a_inf = isinf(a)
      b_inf = isinf(b)
      infs = bitwise_or(a_inf, b_inf)
  
      nans_or_infs = bitwise_or(nans, infs)
  
      # close is now an array of 0's except where elements are not nan or inf
      # and are within the tolerance.
      close = bitwise_and(close_prelim, bitwise_not(nans_or_infs))
  
      # deal with signed inf values. this will make an array inf_eq of 0's
      # except where inf values have the same sign.
      both_infs = bitwise_and(a_inf, b_inf)
      inf_signs_eq = eq(a_inf * sgn(a), b_inf * sgn(b))
      inf_eq = bitwise_and(both_infs, inf_signs_eq)
  
      # now create the potential result combining close and inf_eq
      close_with_infs = bitwise_or(close, inf_eq)
  
      # deal with comparing nan's.
      if equal_nan:
          both_nans = bitwise_and(a_nan, b_nan)
          return bitwise_or(close_with_infs, both_nans)
      # otherwise nan's aren't considered close.
      else:
          return close_with_infs
  
  
  ##########################
  # Bit-wise
  ##########################
  
  
  @scalar_elemwise
  def and_(a, b):
      """bitwise a & b"""
  
  
  bitwise_and = and_  # numpy name for it
  
  
  @scalar_elemwise
  def or_(a, b):
      """bitwise a | b"""
  
  
  bitwise_or = or_  # numpy name for it
  
  
  @scalar_elemwise
  def xor(a, b):
      """bitwise a ^ b"""
  
  
  bitwise_xor = xor  # numpy name for it
  
  
  @scalar_elemwise
  def invert(a):
      """bitwise ~a"""
  
  
  bitwise_not = invert  # numpy alias for it
  
  ##########################
  # Math
  ##########################
  
  
  @scalar_elemwise
  def abs(a):
      """|`a`|"""
  
  
  # These are deprecated and will be removed
  abs_ = abs
  
  
  pprint.assign(abs, printing.PatternPrinter(("|%(0)s|", -1000)))
  
  
  @scalar_elemwise
  def exp(a):
      """e^`a`"""
  
  
  @scalar_elemwise
  def exp2(a):
      """2^`a`"""
  
  
  @scalar_elemwise
  def expm1(a):
      """e^`a` - 1"""
  
  
  @scalar_elemwise
  def neg(a):
      """-a"""
  
  
  @scalar_elemwise
  def reciprocal(a):
      """1.0/a"""
  
  
  # This is deprecated and will be removed
  inv = reciprocal
  
  
  @scalar_elemwise
  def log(a):
      """base e logarithm of a"""
  
  
  @scalar_elemwise
  def log2(a):
      """base 2 logarithm of a"""
  
  
  @scalar_elemwise
  def log10(a):
      """base 10 logarithm of a"""
  
  
  @scalar_elemwise
  def log1p(a):
      """log(1+a)"""
  
  
  @scalar_elemwise
  def sgn(a):
      """sign of a"""
  
  
  @scalar_elemwise
  def ceil(a):
      """ceiling of a"""
  
  
  @scalar_elemwise
  def floor(a):
      """floor of a"""
  
  
  @scalar_elemwise
  def trunc(a):
      """trunc of a"""
  
  
  def iround(a, mode=None):
      """cast(round(a,mode),'int64')"""
      return cast(round(a, mode), "int64")
  
  
  def round(a, mode=None):
      """round_mode(a) with mode in [half_away_from_zero, half_to_even].
      Default to half_to_even."""
      if mode is None:
          mode = "half_to_even"
          if config.warn__round:
              warnings.warn(
                  "aesara.tensor.round() changed its default from"
                  " `half_away_from_zero` to `half_to_even` to have"
                  " the same default as NumPy. Use the Aesara flag"
                  " `warn__round=False` to disable this warning."
              )
      if mode == "half_away_from_zero":
          return round_half_away_from_zero(a)
      elif mode == "half_to_even":
          return round_half_to_even(a)
      else:
          raise Exception(f"round mode {mode} is not implemented.")
  
  
  @scalar_elemwise
  def round_half_to_even(a):
      """round_half_to_even(a)"""
  
  
  @scalar_elemwise
  def round_half_away_from_zero(a):
      """round_half_away_from_zero(a)"""
  
  
  @scalar_elemwise
  def sqr(a):
      """square of a"""
  
  
  # alias to sqr, included to maintain similarity with numpy interface
  square = sqr
  
  
  def cov(m, y=None, rowvar=True, bias=False, ddof=None, fweights=None, aweights=None):
      """Calculate the covariance matrix.
  
      Covariance indicates the level to which two variables vary together.
      If we examine N-dimensional samples, :math:`m = [x_1, x_2, ... x_N]^T`,
      then the covariance matrix element :math:`C_{ij}` is the covariance of
      :math:`x_i` and :math:`x_j`. The element :math:`C_{ii}` is the variance
      of :math:`x_i`. Code and docstring ported from numpy.
  
      Parameters
      ==========
      m : array_like
          A 2-D array containing multiple variables and observations.
          Each row of `m` represents a variable, and each column is
          observations of all those variables.
      y : array_like, optional
          An additional set of variables and observations. `y` has the same form
          as that of `m`.
      rowvar : bool, optional
          If `rowvar` is True (default), then each row represents a
          variable, with observations in the columns. Otherwise, the relationship
          is transposed: each column represents a variable, while the rows
          contain observations.
      bias : bool, optional
          Default normalization (False) is by ``(N - 1)``, where ``N`` is the
          number of observations given (unbiased estimate). If `bias` is True, then
          normalization is by ``N``. These values can be overridden by using the
          keyword ``ddof``.
      ddof : int, optional
          If not ``None`` the default value implied by `bias` is overridden.
          The default value is ``None``.
  
      Returns
      =======
      out : The covariance matrix of the variables.
  
      """
  
      if fweights is not None:
          raise NotImplementedError("fweights are not implemented")
      if aweights is not None:
          raise NotImplementedError("aweights are not implemented")
  
      if not rowvar and m.shape[0] != 1:
          m = m.T
  
      if y is not None:
          if not rowvar and y.shape[0] != 1:
              y = y.T
          m = concatenate((m, y), axis=0)
  
      if ddof is None:
          if not bias:
              ddof = 1
          else:
              ddof = 0
  
      # Determine the normalization
      fact = m.shape[1] - ddof
  
      m -= m.mean(axis=1, keepdims=1)
      c = m.dot(m.T)
      c *= constant(1) / fact
      return c.squeeze()
  
  
  @scalar_elemwise
  def sqrt(a):
      """square root of a"""
  
  
  @scalar_elemwise
  def deg2rad(a):
      """convert degree a to radian"""
  
  
  @scalar_elemwise
  def rad2deg(a):
      """convert radian a to degree"""
  
  
  @scalar_elemwise
  def cos(a):
      """cosine of a"""
  
  
  @scalar_elemwise
  def arccos(a):
      """arccosine of a"""
  
  
  @scalar_elemwise
  def sin(a):
      """sine of a"""
  
  
  @scalar_elemwise
  def arcsin(a):
      """arcsine of a"""
  
  
  @scalar_elemwise
  def tan(a):
      """tangent of a"""
  
  
  @scalar_elemwise
  def arctan(a):
      """arctangent of a"""
  
  
  @scalar_elemwise
  def arctan2(a, b):
      """arctangent of a / b"""
  
  
  @scalar_elemwise
  def cosh(a):
      """hyperbolic cosine of a"""
  
  
  @scalar_elemwise
  def arccosh(a):
      """hyperbolic arc cosine of a"""
  
  
  @scalar_elemwise
  def sinh(a):
      """hyperbolic sine of a"""
  
  
  @scalar_elemwise
  def arcsinh(a):
      """hyperbolic arc sine of a"""
  
  
  @scalar_elemwise
  def tanh(a):
      """hyperbolic tangent of a"""
  
  
  @scalar_elemwise
  def arctanh(a):
      """hyperbolic arc tangent of a"""
  
  
  @scalar_elemwise
  def erf(a):
      """error function"""
  
  
  @scalar_elemwise
  def erfc(a):
      """complementary error function"""
  
  
  @scalar_elemwise
  def erfcx(a):
      """scaled complementary error function"""
  
  
  @scalar_elemwise
  def erfinv(a):
      """inverse error function"""
  
  
  @scalar_elemwise
  def erfcinv(a):
      """inverse complementary error function"""
  
  
  @scalar_elemwise
  def gamma(a):
      """gamma function"""
  
  
  @scalar_elemwise
  def gammaln(a):
      """log gamma function"""
  
  
  @scalar_elemwise
  def psi(a):
      """derivative of log gamma function"""
  
  
  @scalar_elemwise
  def tri_gamma(a):
      """second derivative of the log gamma function"""
  
  
  @scalar_elemwise
  def chi2sf(x, k):
      """chi squared survival function"""
  
  
  @scalar_elemwise
  def gammainc(k, x):
      """Regularized lower gamma function"""
  
  
  @scalar_elemwise
  def gammaincc(k, x):
      """Regularized upper gamma function"""
  
  
  @scalar_elemwise
  def gammau(k, x):
      """Upper incomplete gamma function."""
  
  
  @scalar_elemwise
  def gammal(k, x):
      """Lower incomplete gamma function."""
  
  
  @scalar_elemwise
  def j0(x):
      """Bessel function of the first kind of order 0."""
  
  
  @scalar_elemwise
  def j1(x):
      """Bessel function of the first kind of order 1."""
  
  
  @scalar_elemwise
  def jv(v, x):
      """Bessel function of the first kind of order v (real)."""
  
  
  @scalar_elemwise
  def i0(x):
      """Modified Bessel function of the first kind of order 0."""
  
  
  @scalar_elemwise
  def i1(x):
      """Modified Bessel function of the first kind of order 1."""
  
  
  @scalar_elemwise
  def iv(v, x):
      """Modified Bessel function of the first kind of order v (real)."""
  
  
  @scalar_elemwise
  def sigmoid(x):
      """Logistic sigmoid function (1 / (1 + exp(x)), also known as expit or inverse logit"""
  
  
  expit = sigmoid
  
  
  @scalar_elemwise
  def softplus(x):
      """Compute log(1 + exp(x)), also known as softplus or log1pexp"""
  
  
  log1pexp = softplus
  
  
  @scalar_elemwise
  def log1mexp(x):
      """Compute log(1 - exp(x)), also known as log1mexp"""
  
  
  @scalar_elemwise
  def betainc(a, b, x):
      """Regularized incomplete beta function"""
  
  
  @scalar_elemwise
  def real(z):
      """Return real component of complex-valued tensor `z`"""
  
  
  _tensor_py_operators.real = property(real)
  
  
  @scalar_elemwise
  def imag(z):
      """Return imaginary component of complex-valued tensor `z`"""
  
  
  _tensor_py_operators.imag = property(imag)
  
  
  @scalar_elemwise
  def angle(z):
      """Return polar-coordinate angle of complex-valued tensor `z`"""
  
  
  @scalar_elemwise  # numpy.complex cannot build tensors
  def complex(real, imag):
      """Return complex-valued tensor with `real` and `imag` components"""
  
  
  @scalar_elemwise
  def conj(z):
      """Return the complex conjugate of `z`."""
  
  
  @scalar_elemwise
  def complex_from_polar(abs, angle):
      """Return complex-valued tensor from polar coordinate specification."""
  
  
  class Mean(CAReduce):
      def __init__(self, axis=None):
          super().__init__(aes.add, axis)
          assert self.axis is None or len(self.axis) == 1
  
      def __str__(self):
          if self.axis is not None:
              return "Mean{%s}" % (", ".join(str(x) for x in self.axis))
          else:
              return "Mean"
  
      def _output_dtype(self, idtype):
          # we want to protect against overflow
          return "float64"
  
      def perform(self, node, inp, out):
          (input,) = inp
          (output,) = out
          if self.axis is None:
              axis = None
          else:
              axis = self.axis[0]
          # numpy.asarray is needed as otherwise we can end up with a
          # numpy scalar.
          output[0] = np.asarray(np.mean(input, dtype="float64", axis=axis))
  
      def c_code(self, node, name, inames, onames, sub):
  
          ret = super().c_code(node, name, inames, onames, sub)
  
          if self.axis is not None:
              return ret
  
          # TODO: c_code perform support only axis is None
          return (
              ret
              + f"""
    *((double *)PyArray_DATA({onames[0]})) /= PyArray_SIZE({inames[0]});
    """
          )
  
  
  # TODO: implement the grad. When done and tested, you can make this the default
  # version.
  #    def grad(self, (x,), (gout,)):
  #      import pdb;pdb.set_trace()
  #      return grad(mean(x, self.axis, op=False),[x])
  
  
  def mean(input, axis=None, dtype=None, op=False, keepdims=False, acc_dtype=None):
      """
      Computes the mean value along the given axis(es) of a tensor `input`.
  
      Parameters
      ----------
      axis : None or int or (list of int) (see `Sum`)
          Compute the mean along this axis of the tensor.
          None means all axes (like numpy).
      dtype: None or string
          Dtype to cast the result of the inner summation into.
          For instance, by default, a sum of a float32 tensor will be
          done in float64 (acc_dtype would be float64 by default),
          but that result will be casted back in float32.
      keepdims: bool
          If this is set to True, the axes which are reduced are
          left in the result as dimensions with size one. With this option,
          the result will broadcast correctly against the original tensor.
      acc_dtype: None or string
          Dtype to use for the inner summation. This will not
          necessarily be the dtype of the output (in particular
          if it is a discrete (int/uint) dtype, the output will
          be in a float type). If None, then we use the same rules as `sum()`.
  
      Notes
      -----
      For gpu, if you specify dtype=float32, everything will be done on the gpu.
  
      """
      input = as_tensor_variable(input)
      if op:
          if dtype not in (None, "float64"):
              raise NotImplementedError(
                  "The Mean op does not support the dtype argument, "
                  "and will always use float64. If you want to specify "
                  "the dtype, call tensor.mean(..., op=False).",
                  dtype,
              )
          if acc_dtype not in (None, "float64"):
              raise NotImplementedError(
                  "The Mean op does not support the acc_dtype argument, "
                  "and will always use float64. If you want to specify "
                  "acc_dtype, call tensor.mean(..., op=False).",
                  dtype,
              )
          out = Mean(axis)(input)
          if keepdims:
              out = makeKeepDims(input, out, axis)
          return out
  
      if dtype is not None:
          # The summation will be done with the specified dtype.
          # sum() will complain if it is not suitable.
          sum_dtype = dtype
      else:
          sum_dtype = None
          # float16 overflows on the cast way too often
          if input.dtype == "float16":
              sum_dtype = "float32"
  
      s = sum(input, axis=axis, dtype=sum_dtype, keepdims=keepdims, acc_dtype=acc_dtype)
      shp = shape(input)
  
      # Cast shp into a float type
      # TODO Once we have a consistent casting policy, we could simply
      # use true_div.
      if s.dtype in ("float16", "float32", "complex64"):
          shp = cast(shp, "float32")
      else:
          shp = cast(shp, "float64")
  
      if axis is None:
          axis = list(range(input.ndim))
      elif isinstance(axis, (int, np.integer)):
          axis = [axis]
      elif isinstance(axis, np.ndarray) and axis.ndim == 0:
          axis = [int(axis)]
      else:
          axis = [int(a) for a in axis]
  
      # This sequential division will possibly be optimized by Aesara:
      for i in axis:
          s = true_div(s, shp[i])
  
      # This can happen when axis is an empty list/tuple
      if s.dtype != shp.dtype and s.dtype in discrete_dtypes:
          s = cast(s, shp.dtype)
  
      if dtype == "float16" or (dtype is None and input.dtype == "float16"):
          s = cast(s, "float16")
      s.name = "mean"
      return s
  
  
  def var(input, axis=None, ddof=0, keepdims=False, corrected=False):
      """
      Computes the variance along the given axis(es) of a tensor `input`.
  
      Parameters
      ----------
      axis: None or int or (list of int) (see `Sum`)
          Compute the variance along this axis of the tensor.
          None means all axes (like numpy).
      ddof: Degrees of freedom; 0 would compute the ML estimate, 1 would compute
          the unbiased estimate.
      keepdims : bool
          If this is set to True, the axes which are reduced are
          left in the result as dimensions with size one. With this option,
          the result will broadcast correctly against the original tensor.
      corrected : bool
          If this is set to True, the 'corrected_two_pass' algorithm is
          used to compute the variance.
          Refer : http://www.cs.yale.edu/publications/techreports/tr222.pdf
  
      Notes
      -----
      Default uses the two-pass algorithm (reference below).
      https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Two-pass_algorithm
      Also supports 'corrected_two_pass' algorithm (using the 'corrected' flag)
      which is numerically more stable. There exist other implementations that
      offer better stability, but probably slower.
  
      """
  
      if isinstance(ddof, (bool)):
          raise ValueError(
              "Parameter keepdims is now at index 3: (input, \
                            axis=None, ddof=0, keepdims=False, corrected=False)"
          )
  
      input_ndim = input.type.ndim
      if axis is None:
          axis = list(range(input_ndim))
      elif isinstance(axis, (int, np.integer)):
          axis = [axis]
      elif isinstance(axis, np.ndarray) and axis.ndim == 0:
          axis = [int(axis)]
      else:
          axis = [int(a) for a in axis]
  
      # compute the axis-wise mean
      mean_input = mean(input, axis, keepdims=True)
  
      # center the input
      centered_input = input - mean_input
  
      # return the mean sqr
      two = constant(2, dtype=centered_input.dtype)
      if ddof == 0:
          v = mean((centered_input ** two), axis, keepdims=keepdims)
      else:
          shp = shape(input) - ddof
          v = sum((centered_input ** two), axis=axis, keepdims=keepdims)
          for i in axis:
              v = true_div(v, shp[i])
  
      # use 'corrected_two_pass' algorithm
      if corrected:
          if ddof == 0:
              error = mean(centered_input, axis, keepdims=keepdims) ** 2
          else:
              shp = shape(input) - ddof
              shp_inp = shape(input)
              error = sum(centered_input, axis=axis, keepdims=keepdims) ** 2
              for i in axis:
                  error = true_div(error, shp[i] * shp_inp[i])
          v = v - error
  
      v.name = "var"
      return v
  
  
  def std(input, axis=None, ddof=0, keepdims=False, corrected=False):
      """
      Computes the standard deviation along the given axis(es) of a tensor `input`.
  
      Parameters
      ----------
      axis: None or int or (list of int) (see `Sum`)
          Compute the variance along this axis of the tensor.
          None means all axes (like numpy).
      ddof: Degrees of freedom; 0 would compute the ML estimate, 1 would compute
          the unbiased estimate.
      keepdims : bool
          If this is set to True, the axes which are reduced are
          left in the result as dimensions with size one. With this option,
          the result will broadcast correctly against the original tensor.
      corrected : bool
          If this is set to True, the 'corrected_two_pass' algorithm is
          used to compute the variance.
          Refer : http://www.cs.yale.edu/publications/techreports/tr222.pdf
  
      Notes
      -----
      It calls 'var()' and 'var()' uses the two-pass algorithm (reference below).
      https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Two-pass_algorithm
      Function 'var()' also supports 'corrected_two_pass' algorithm (using the
      'corrected' flag) which is numerically more stable. There exist other
      implementations that offer better stability, but probably slower.
  
      """
  
      if isinstance(ddof, (bool)):
          raise ValueError(
              "Parameter keepdims is now at index 3: (input, \
                            axis=None, ddof=0, keepdims=False, corrected=False)"
          )
  
      ret = sqrt(
          var(input=input, axis=axis, ddof=ddof, keepdims=keepdims, corrected=corrected)
      )
      ret.name = "std"
      return ret
  
  
  @scalar_elemwise(symbolname="scalar_maximum")
  def maximum(x, y):
      """elemwise maximum. See max for the maximum in one tensor"""
      # see decorator for function body
  
  
  @scalar_elemwise(symbolname="scalar_minimum")
  def minimum(x, y):
      """elemwise minimum. See min for the minimum in one tensor"""
      # see decorator for function body
  
  
  def divmod(x, y):
      """elementvise divmod, using floor_div and mod_check"""
      return floor_div(x, y), mod_check(x, y)
  
  
  @scalar_elemwise
  def add(a, *other_terms):
      """elementwise addition"""
      # see decorator for function body
  
  
  @scalar_elemwise
  def sub(a, b):
      """elementwise subtraction"""
      # see decorator for function body
  
  
  @scalar_elemwise
  def mul(a, *other_terms):
      """elementwise multiplication"""
      # see decorator for function body
  
  
  @scalar_elemwise
  def true_div(a, b):
      """elementwise [true] division (inverse of multiplication)"""
      # see decorator for function body
  
  
  @scalar_elemwise
  def int_div(a, b):
      """elementwise [floor] division (inverse of multiplication)"""
      # see decorator for function body
  
  
  # floor_div and int_div are the same thing
  floor_div = int_div
  
  
  def ceil_intdiv(a, b):
      """
      Safely compute ceil(float_division(a, b)).
  
      Works for all dtypes, but mostly useful when a and b are int.
  
      """
      # If a and b are int with not many significant bits, we could
      # cast them to float to avoid doing the modulo. We do not know if this
      # is faster or not. But this is not safe for int64 as the cast will
      # lose precision.
      # e.g.: cast(cast(a, scalar.upcast(a, 'float32')) / b, aes.upcast(a, b))
  
      # We cast for the case when a and b are uint*. Otherwise neq will
      # force their upcast to int.
      div = int_div(a, b)
      ret = cast(neq(a % b, 0), div.dtype) + div
      assert ret.dtype == aes.upcast(div.owner.inputs[0], div.owner.inputs[1])
      return ret
  
  
  def mod_check(x, y):
      """Make sure we do not try to use complex numbers."""
      if (
          as_tensor_variable(x).dtype in complex_dtypes
          or as_tensor_variable(y).dtype in complex_dtypes
      ):
          # Currently forbidden.
          raise aes.Mod.complex_error
      else:
          return mod(x, y)
  
  
  @scalar_elemwise
  def mod(a, b):
      """elementwise modulo"""
      # see decorator for function body
  
  
  @scalar_elemwise
  def pow(a, b):
      """elementwise power"""
      # see decorator for function body
  
  
  @scalar_elemwise
  def clip(x, min, max):
      """
      Clip x to be between min and max.
  
      Note that when `x` is equal to the boundaries, the output is considered
      to be `x`, so at these points, the gradient of the cost wrt the output
      will be propagated to `x`, not to `min` nor `max`. In other words,
      on these points, the gradient wrt `x` will be equal to the gradient wrt
      the output, and the gradient wrt `min` and `max` will be zero.
  
      """
      # see decorator for function body
      # for grep: clamp, bound
  
  
  pprint.assign(add, printing.OperatorPrinter("+", -2, "either"))
  pprint.assign(mul, printing.OperatorPrinter("*", -1, "either"))
  pprint.assign(sub, printing.OperatorPrinter("-", -2, "left"))
  pprint.assign(neg, printing.OperatorPrinter("-", 0, "either"))
  pprint.assign(true_div, printing.OperatorPrinter("/", -1, "left"))
  pprint.assign(int_div, printing.OperatorPrinter("//", -1, "left"))
  pprint.assign(pow, printing.OperatorPrinter("**", 1, "right"))
  
  
  class Dot(Op):
      """
      Computes the dot product of two variables. For two matrices, this is
      equivalent to matrix multiplication. For two vectors, this is the inner
      product.
  
      Notes
      -----
      Matrix-matrix products are sometimes optimized to Dot22 or Gemm ops
      (see tensor.blas).
      Vector-vector products are sometimes optimized to Ger or CGer (see
      tensor.blas).
      Matrix-vector products are sometimes optimized to Gemv, CGemv (see
      tensor.blas).
  
      """
  
      __props__ = ()
  
      # the rationale for Dot22 is related to getting GEMM Ops into the
      # graph.  See Dot22 in tensor.blas for details.
  
      def make_node(self, *inputs):
          inputs = list(map(as_tensor_variable, inputs))
  
          if len(inputs) != 2:
              raise TypeError(f"Two arguments required, {len(inputs)} given ")
          if inputs[0].ndim not in (1, 2):
              raise TypeError(
                  "Input 0 (0-indexed) must have ndim of "
                  f"1 or 2, {int(inputs[0].ndim)} given. Consider calling "
                  "aesara.tensor.dot instead."
              )
          if inputs[1].ndim not in (1, 2):
              raise TypeError(
                  "Input 1 (0-indexed) must have ndim of "
                  f"1 or 2, {int(inputs[1].ndim)} given. Consider calling "
                  "aesara.tensor.dot instead."
              )
  
          i_broadcastables = [input.type.broadcastable for input in inputs]
          bx, by = i_broadcastables
          if len(by) == 2:  # y is a matrix
              bz = bx[:-1] + by[-1:]
          elif len(by) == 1:  # y is vector
              bz = bx[:-1]
  
          i_dtypes = [input.type.dtype for input in inputs]
          outputs = [tensor(aes.upcast(*i_dtypes), bz)]
          return Apply(self, inputs, outputs)
  
      def perform(self, node, inp, out):
          x, y = inp
          (z,) = out
  
          # the asarray is here because dot between two vectors
          # gives a numpy float object but we need to return a 0d
          # ndarray
          z[0] = np.asarray(np.dot(x, y))
  
      def grad(self, inp, grads):
  
          x, y = inp
          (gz,) = grads
          xdim, ydim, gdim = x.type.ndim, y.type.ndim, gz.type.ndim
  
          # grad is scalar, so x is vector and y is vector
          if gdim == 0:
              xgrad = gz * y
              ygrad = gz * x
  
          # x is vector, y is matrix, grad is vector
          elif xdim == 1 and ydim == 2:
              xgrad = dot(gz, y.T)
              ygrad = outer(x.T, gz)
  
          # x is matrix, y is vector, grad is vector
          elif xdim == 2 and ydim == 1:
              xgrad = outer(gz, y.T)
              ygrad = dot(x.T, gz)
  
          # x is matrix, y is matrix, grad is matrix
          elif xdim == ydim == 2:
              xgrad = dot(gz, y.T)
              ygrad = dot(x.T, gz)
  
          # If x or y contain broadcastable dimensions but only one of
          # them know that a matching dimensions is broadcastable, the
          # above code don't always return the right broadcast pattern.
          # This cause problem down the road. See gh-1461.
          if xgrad.broadcastable != x.broadcastable:
              xgrad = patternbroadcast(xgrad, x.broadcastable)
          if ygrad.broadcastable != y.broadcastable:
              ygrad = patternbroadcast(ygrad, y.broadcastable)
  
          rval = xgrad, ygrad
  
          for elem in rval:
              assert elem.dtype.find("float") != -1
  
          return rval
  
      def R_op(self, inputs, eval_points):
          # R_op for a \dot b evaluated at c for a and d for b is
          # simply c \dot b + a \dot d
  
          assert len(inputs) == 2
          assert len(eval_points) == 2
          if eval_points[0] is None and eval_points[1] is None:
              return [None]
  
          if eval_points[0]:
              t1 = self(eval_points[0], inputs[1])
          if eval_points[1]:
              t2 = self(inputs[0], eval_points[1])
  
          if eval_points[0] and eval_points[1]:
              return [t1 + t2]
          elif eval_points[0]:
              return [t1]
          else:
              return [t2]
  
      def infer_shape(self, fgraph, node, shapes):
          xshp, yshp = shapes
          x, y = node.inputs
  
          # vector / vector
          if x.ndim == 1 and y.ndim == 1:
              return [()]
          # matrix / vector
          if x.ndim == 2 and y.ndim == 1:
              return [xshp[:-1]]
          # vector / matrix
          if x.ndim == 1 and y.ndim == 2:
              return [yshp[-1:]]
          # matrix / matrix
          if x.ndim == 2 and y.ndim == 2:
              return [xshp[:-1] + yshp[-1:]]
          raise NotImplementedError()
  
      def __str__(self):
          return "dot"
  
  
  _dot = Dot()
  pprint.assign(
      _dot, printing.OperatorPrinter(printing.special["middle_dot"], -1, "left")
  )
  
  
  def dot(l, r):
      """Return a symbolic dot product.
  
      This is designed to work with both sparse and dense tensors types.
      """
  
      if not isinstance(l, Variable):
          l = as_tensor_variable(l)
  
      if not isinstance(r, Variable):
          r = as_tensor_variable(r)
  
      try:
          res = l.__dot__(r)
          if res is NotImplemented:
              raise NotImplementedError
      except (NotImplementedError, AttributeError, TypeError):
          res = r.__rdot__(l)
          if res is NotImplemented:
              raise NotImplementedError()
  
      return res
  
  
  def dense_dot(a, b):
      """
      Computes the dot product of two variables.
  
      For two matrices, this is equivalent to matrix multiplication.
      For two vectors, this is the inner product.
      When one variable is a scalar, this is like elementwise multiplication.
      For N dimensions, this is a sum product over the last axis
      of the first array and the second-to-last axis of the second array:
  
          dot(a, b)[i,j,k,m] = sum(a[i,j,:] * b[k,:,m])
  
      Note that this dot function does one of three things, in the following
      sequence:
  
          1.  If either a or b is scalar, it returns the elementwise product
              without calling the Aesara Dot op.
  
          2.  If either a or b has more than 2 dimensions, it calls Aesara's
              tensordot function with appropriate axes. The tensordot function
              expresses high-dimensional dot products in terms of 2D matrix
              multiplications, so it may be possible to further optimize for
              performance.
  
          3.  If both a and b have either 1 or 2 dimensions, it calls Aesara's
              Dot op on a and b.
  
      Notes
      -----
      Matrix-matrix products are sometimes optimized to Dot22 or Gemm ops
      (see tensor.blas).
      Vector-vector products are sometimes optimized to Ger or CGer (see
      tensor.blas).
      Matrix-vector products are sometimes optimized to Gemv, CGemv (see
      tensor.blas).
  
      """
      a, b = as_tensor_variable(a), as_tensor_variable(b)
  
      if a.ndim == 0 or b.ndim == 0:
          return a * b
      elif a.ndim > 2 or b.ndim > 2:
          return tensordot(a, b, [[a.ndim - 1], [np.maximum(0, b.ndim - 2)]])
      else:
          return _dot(a, b)
  
  
  def _tensordot_as_dot(a, b, axes, dot, batched):
      """
      Reduces a tensor dot product to a matrix or vector dot product. Based
      on code from Tijmen Tieleman's gnumpy
      (http://www.cs.toronto.edu/~tijmen/gnumpy.html).
  
      Please see the documentation of tensordot for the meaning of the a, b
      and axes arguments.
  
      :param dot: a function that accepts two symbolic variables and computes
                  the appropriate dot product (e.g. dot, batched_dot)
      :type dot: function
  
      :param batched: whether to treat the first axis of a and b as a batch
                      axis.  If so, this axis will be preserved in the output,
                      allowing this function to be used also for batched
                      tensor dot products.
      :type batched: boolean
  
      :returns: a tensor with shape equal to the concatenation of a's shape
                (less any dimensions that were summed over) and b's shape
                (less the first dimension and any dimensions that were summed
                over).
      :rtype: symbolic tensor
      """
      a, b = as_tensor_variable(a), as_tensor_variable(b)
  
      if not np.isscalar(axes) and len(axes) != 2:
          raise ValueError(
              "Axes should be an integer or a "
              "list/tuple of len 2 ({axes} was provided)"
          )
  
      # if 'axes' is a number of axes to multiply and sum over (trailing axes
      # of a, leading axes of b), we can just reshape and use dot.
      elif np.isscalar(axes):
          axes = int(axes)
  
          for operand_name, operand in (("a", a), ("b", b)):
              if axes > operand.ndim:
                  raise ValueError(
                      f"axes can not be larger than the dimension of {operand_name} "
                      f"({operand_name}.ndim={operand.ndim}, axes={axes})"
                  )
              if batched and axes == operand.ndim:
                  raise ValueError(
                      "axes to sum over must not include the batch axis "
                      f"of {operand_name} ({operand_name}.ndim={operand.ndim}, axes={axes})"
                  )
  
          batch_axes = 1 if batched else 0
          a_outaxes = slice(0, a.ndim - axes)
          b_outaxes = slice(batch_axes + axes, b.ndim)
          outshape = concatenate([a.shape[a_outaxes], b.shape[b_outaxes]])
          outbcast = a.broadcastable[a_outaxes] + b.broadcastable[b_outaxes]
          outndim = len(outbcast)
  
          a_shape = [1] * 2
          b_shape = [1] * 2
  
          # compute total size of summed axes
          for i in range(0, axes):
              a_shape[1] *= a.shape[-(i + 1)]
              b_shape[0] *= b.shape[batch_axes + i]
          # compute total size of other axes
          for i in range(0, a.ndim - axes - batch_axes):
              a_shape[0] *= a.shape[batch_axes + i]
          for i in range(0, b.ndim - axes - batch_axes):
              b_shape[1] *= b.shape[-(i + 1)]
  
          if batched:
              a_shape.insert(0, a.shape[0])
              b_shape.insert(0, b.shape[0])
  
          a_reshaped = a.reshape(a_shape)
          b_reshaped = b.reshape(b_shape)
  
          out_reshaped = dot(a_reshaped, b_reshaped)
          out = out_reshaped.reshape(outshape, outndim)
          # Make sure the broadcastable pattern of the result is correct,
          # since some shape information can be lost in the reshapes.
          return patternbroadcast(out, outbcast)
  
      # if 'axes' is a list, transpose a and b such that the summed axes of a
      # are last and the summed axes of b are first.
      else:
          axes = [as_list(axes_) for axes_ in axes]
  
          if len(axes[0]) != len(axes[1]):
              raise ValueError("Axes elements must have the same length.")
  
          for i, (operand_name, operand) in enumerate((("a", a), ("b", b))):
              if len(axes[i]) > operand.ndim:
                  raise ValueError(
                      f"axes[{i}] should be array_like with length less than "
                      f"the dimensions of {operand_name} ({operand_name}.ndim={operand.ndim}, len(axes[0])={len(axes[i])})."
                  )
              if len(axes[i]) > 0 and np.max(axes[i]) >= operand.ndim:
                  raise ValueError(
                      f"axes[{i}] contains dimensions greater than or equal "
                      f"to {operand_name}.ndim ({operand_name}.ndim={operand.ndim}, max(axes[0])={np.max(np.array(axes[i]))})."
                  )
              if batched and 0 in axes[i]:
                  raise ValueError(
                      "axes to sum over must not contain the batch axis "
                      f"(axes[{i}]={axes[i]})"
                  )
  
          batch_axes = [0] if batched else []
          other_axes = [
              [x for x in range(operand.ndim) if x not in axes[i] and x not in batch_axes]
              for i, operand in enumerate((a, b))
          ]
  
          a_shuffled = a.dimshuffle(batch_axes + other_axes[0] + axes[0])
          b_shuffled = b.dimshuffle(batch_axes + axes[1] + other_axes[1])
  
          # now that a and b are in the right order, recur with integer axes
          return _tensordot_as_dot(
              a_shuffled, b_shuffled, len(axes[0]), dot=dot, batched=batched
          )
  
  
  def tensordot(a, b, axes=2):
      """
      Compute a generalized dot product over provided axes.
  
      Given two tensors a and b, tensordot computes a generalized dot product over
      the provided axes. Aesara's implementation reduces all expressions to
      matrix or vector dot products and is based on code from Tijmen Tieleman's
      gnumpy (http://www.cs.toronto.edu/~tijmen/gnumpy.html).
  
      Parameters
      ----------
      a: symbolic tensor
          The first tensor variable.
      b: symbolic tensor
          The second tensor variable
      axes: int or array-like of length 2
          If an integer, the number of axes to sum over.
          If an array, it must have two array elements containing the axes
          to sum over in each tensor.
  
          Note that the default value of 2 is not guaranteed to work
          for all values of a and b, and an error will be raised if
          that is the case. The reason for keeping the default is to
          maintain the same signature as numpy's tensordot function
          (and np.tensordot raises analogous errors for non-compatible
          inputs).
  
          If an integer i, it is converted to an array containing
          the last i dimensions of the first tensor and the first
          i dimensions of the second tensor:
              axes = [list(range(a.ndim - i, b.ndim)), list(range(i))]
  
          If an array, its two elements must contain compatible axes
          of the two tensors. For example, [[1, 2], [2, 0]] means sum
          over the 2nd and 3rd axes of a and the 3rd and 1st axes of b.
          (Remember axes are zero-indexed!) The 2nd axis of a and the
          3rd axis of b must have the same shape; the same is true for
          the 3rd axis of a and the 1st axis of b.
  
      Returns
      -------
      symbolic tensor
          A tensor with shape equal to the concatenation of a's shape
          (less any dimensions that were summed over) and b's shape
          (less any dimensions that were summed over).
  
      Examples
      --------
      It may be helpful to consider an example to see what tensordot does.
      Aesara's implementation is identical to NumPy's. Here a has shape (2, 3, 4)
      and b has shape (5, 6, 4, 3). The axes to sum over are [[1, 2], [3, 2]] --
      note that a.shape[1] == b.shape[3] and a.shape[2] == b.shape[2]; these axes
      are compatible. The resulting tensor will have shape (2, 5, 6) -- the
      dimensions that are not being summed:
  
      >>> a = np.random.random((2,3,4))
      >>> b = np.random.random((5,6,4,3))
  
      #tensordot
      >>> c = np.tensordot(a, b, [[1,2],[3,2]])
  
      #loop replicating tensordot
      >>> a0, a1, a2 = a.shape
      >>> b0, b1, _, _ = b.shape
      >>> cloop = np.zeros((a0,b0,b1))
  
      #loop over non-summed indices -- these exist
      #in the tensor product.
      >>> for i in range(a0):
      ...     for j in range(b0):
      ...         for k in range(b1):
      ...             #loop over summed indices -- these don't exist
      ...             #in the tensor product.
      ...             for l in range(a1):
      ...                 for m in range(a2):
      ...                     cloop[i,j,k] += a[i,l,m] * b[j,k,m,l]
  
      >>> np.allclose(c, cloop)
      true
  
      This specific implementation avoids a loop by transposing a and b such that
      the summed axes of a are last and the summed axes of b are first. The
      resulting arrays are reshaped to 2 dimensions (or left as vectors, if
      appropriate) and a matrix or vector dot product is taken. The result is
      reshaped back to the required output dimensions.
  
      In an extreme case, no axes may be specified. The resulting tensor
      will have shape equal to the concatenation of the shapes of a and b:
  
      >>> c = np.tensordot(a, b, 0)
      >>> print(a.shape)
      (2,3,4)
      >>> print(b.shape)
      (5,6,4,3)
      >>> print(c.shape)
      (2,3,4,5,6,4,3)
  
      See the documentation of numpy.tensordot for more examples.
  
      """
      return _tensordot_as_dot(a, b, axes, dot=dot, batched=False)
  
  
  def outer(x, y):
      """Return vector-vector outer product.
  
      If an input isn't a vector, we flatten it first.
  
      """
      if x.ndim != 1:
          x = x.flatten()
      if y.ndim != 1:
          y = y.flatten()
      return dot(x.dimshuffle(0, "x"), y.dimshuffle("x", 0))
  
  
  class All(CAReduce):
      """Applies `logical and` to all the values of a tensor along the
      specified axis(es).
  
      """
  
      __props__ = ("axis",)
      nfunc_spec = ("all", 1, 1)
  
      def __init__(self, axis=None):
          super().__init__(aes.and_, axis)
  
      def _output_dtype(self, idtype):
          return "bool"
  
      def __str__(self):
          if self.axis is None:
              return "All"
          else:
              return "All{%s}" % ", ".join(map(str, self.axis))
  
      def make_node(self, input):
          input = as_tensor_variable(input)
          if input.dtype != "bool":
              input = neq(input, 0)
          ret = super().make_node(input)
          return ret
  
      def grad(self, inp, grads):
          (x,) = inp
          return [x.zeros_like(config.floatX)]
  
  
  class Any(CAReduce):
      """Applies `bitwise or` to all the values of a tensor along the
      specified axis(es).
  
      """
  
      __props__ = ("axis",)
      nfunc_spec = ("any", 1, 1)
  
      def __init__(self, axis=None):
          super().__init__(aes.or_, axis)
  
      def _output_dtype(self, idtype):
          return "bool"
  
      def __str__(self):
          if self.axis is None:
              return "Any"
          else:
              return "Any{%s}" % ", ".join(map(str, self.axis))
  
      def make_node(self, input):
          input = as_tensor_variable(input)
          if input.dtype != "bool":
              input = neq(input, 0)
          ret = super().make_node(input)
          return ret
  
      def grad(self, inp, grads):
          (x,) = inp
          return [x.zeros_like(config.floatX)]
  
  
  class Sum(CAReduceDtype):
      """
      Sums all the values of a tensor along the specified axis(es).
  
      Equivalent to `CAReduceDtype(scalar.add, axis=axis, dtype=dtype)`,
      with the difference that this defines the gradient of sum wrt its
      tensor input.
  
      Parameters
      ----------
      axis
          Axis(es) along which the tensor should be summed
          (use None to sum over all axes, and a list or tuple to sum along more
          than one axis).
  
      dtype
          The dtype of the internal accumulator and returned
          tensor. If None, then we use the default dtype which is the same as the
          input tensor's dtype except when:
          - the input dtype is a signed integer of precision < 64 bit, in
          which case we use int64
          - the input dtype is an unsigned integer of precision < 64 bit, in
          which case we use uint64
          This value does not depend on the value of "acc_dtype".
  
      acc_dtype
          The dtype of the internal accumulator.
          If None (default), we use the dtype in the list below,
          or the input dtype if its precision is higher:
          - for int dtypes, we use at least int64;
          - for uint dtypes, we use at least uint64;
          - for float dtypes, we use at least float64;
          - for complex dtypes, we use at least complex128.
  
      """
  
      __props__ = ("axis", "dtype", "acc_dtype")
      nfunc_spec = ("sum", 1, 1)
  
      def __init__(self, axis=None, dtype=None, acc_dtype=None):
          super().__init__(aes.add, axis=axis, dtype=dtype, acc_dtype=acc_dtype)
  
      def __str__(self):
          name = self.__class__.__name__
          axis = ""
          if self.axis is not None:
              axis = ", ".join(str(x) for x in self.axis)
              axis = f"axis=[{axis}], "
          return f"{name}{{{axis}acc_dtype={self.acc_dtype}}}"
  
      def L_op(self, inp, out, grads):
          (x,) = inp
  
          if out[0].dtype not in continuous_dtypes:
              return [x.zeros_like(dtype=config.floatX)]
  
          (gz,) = grads
          gz = as_tensor_variable(gz)
          axis = self.axis
          if axis is None:
              axis = list(range(x.type.ndim))
          if axis == ():
              return (gz,)
          new_dims = []
          i = 0
          for j, _ in enumerate(x.type.broadcastable):
              if j in axis:
                  new_dims.append("x")
              else:
                  new_dims.append(i)
                  i += 1
          ds_op = DimShuffle(gz.type.broadcastable, new_dims)
          gx = Elemwise(aes.second)(x, ds_op(gz))
          return [gx]
  
      def R_op(self, inputs, eval_points):
          # There is just one element in inputs and eval_points, the axis are
          # part of self
          if None in eval_points:
              return [None]
          return self(*eval_points, **dict(return_list=True))
  
  
  def sum(input, axis=None, dtype=None, keepdims=False, acc_dtype=None):
      """
      Computes the sum along the given axis(es) of a tensor `input`.
  
      When axis is None (the default value), the sum is performed
      over the flattened tensor.
  
      For full documentation see `Sum`.
      In particular please pay attention to the important warning when using
      a custom acc_dtype.
  
      Parameters
      ----------
      keepdims: bool
          If this is set to True, the axes which are reduced are left in
          the result as dimensions with size one. With this option, the result
          will broadcast correctly against the original tensor.
  
      """
  
      out = Sum(axis=axis, dtype=dtype, acc_dtype=acc_dtype)(input)
  
      if keepdims:
          out = makeKeepDims(input, out, axis)
      return out
  
  
  pprint.assign(Sum(), printing.FunctionPrinter("sum"))
  
  
  class Prod(CAReduceDtype):
      """
      Multiplies all the values of a tensor along the specified axis(es).
  
      Equivalent to `CAReduce(scalar.mul, axis = axis)`, with the
      difference that this defines the gradient of prod wrt its tensor
      input.
  
      """
  
      __props__ = ("axis", "dtype", "acc_dtype")
      nfunc_spec = ("prod", 1, 1)
  
      def __init__(self, axis=None, dtype=None, acc_dtype=None, no_zeros_in_input=False):
          super().__init__(aes.mul, axis=axis, dtype=dtype, acc_dtype=acc_dtype)
          self.no_zeros_in_input = no_zeros_in_input
  
      def __setstate__(self, dct):
          super().__setstate__(dct)
          # Add default value to be able to reload old pickled objects.
          if "no_zeros_in_input" not in dct:
              self.no_zeros_in_input = False
  
      def L_op(self, inp, out, grads):
          """
          The grad of this Op could be very easy, if it is was not for the case
          where zeros are present in a given "group" (ie. elements reduced
          together to form the product).
  
          If no zeros are found in the elements of the product, then the
          partial derivative of the product relative to one of the elements
          (one of the inputs) is simply the product of the other elements.
          That's easy to see from the chain rule.
  
          Now the trick (with no zeros) is to take the overall product, then
          for every original element, the partial derivative is given by
          this product divided by the element itself (which equals the product
          of the other terms). This is easy to do by broadcasting the original
          product.
  
          (Note that we also need to broadcast-multiply by the
          "incoming gradient", ie. the gradient of the cost relative to the
          output/product).
  
          With zeros, things get more complicated. For a given group, we have 3
          cases:
  
          * No zeros in the group. Use previous trick.
          * If only one zero is present, then the gradient for that element is
              non-zero, but is zero for all others.
          * If more than one zero is present, then all the derivatives are zero.
  
          For the last two cases (with 1 or more zeros), we can't use the
          division trick, as this gives divisions by 0.
  
          Implementing that case-by-case logic is not as trivial, so a bunch of
          hacks are piled down here to do it. Notably, for the "only one zero"
          case, there's a special Op that computes the product of the elements
          in the group, minus the zero (see `ProdWithoutZeros`). The trick is then
          to use the division trick for groups with no zero, to use the
          `ProdWithoutZeros` op where there's only one zero, and to output a
          derivative of zero for any element part of a group with more than
          one zero.
  
          I do this by first counting the number of zeros in each group (see the
          `aet.eq` bits), then taking this or that behavior (see `aet.switch`)
          based on the result of this count.
  
          """
          (prod_in,) = inp
          (gz,) = grads
  
          if out[0].dtype in discrete_dtypes or self.acc_dtype in discrete_dtypes:
              # There is an int conversion in the way
              return [prod_in.zeros_like(dtype=config.floatX)]
  
          # Prepare the broadcasting that is used everywhere to broadcast
          # over the original groups (ie. broadcast over the elements of a given
          # product)
          gz = as_tensor_variable(gz)
          axis = self.axis
          if axis is None:
              axis = list(range(prod_in.type.ndim))
          if axis == ():
              return (gz,)
          new_dims = []
          i = 0
          for j, _ in enumerate(prod_in.type.broadcastable):
              if j in axis:
                  new_dims.append("x")
              else:
                  new_dims.append(i)
                  i += 1
  
          # result of the product, broadcastable over groups
          prod_out = self(prod_in).dimshuffle(new_dims)
          # incoming gradient, broadcastable over groups
          gz = gz.dimshuffle(new_dims)
  
          # division trick if we don't have zeros. This will contain
          # NaNs to be eliminated in the `aet.switch` if we do have zeros.
          grad_case_without_zeros = gz * prod_out / prod_in
  
          if self.no_zeros_in_input:
              # this handles inputs with zeros, but only certain input shapes
              return [grad_case_without_zeros]
          else:
  
              where_zeros = eq(prod_in, 0.0)
              sum_where_zeros = sum(where_zeros, axis=self.axis)
              groups_with_single_zero = eq(sum_where_zeros, 1).dimshuffle(new_dims)
              # tensor with 0 everywhere except for those places where
              # a 0 part of a group with a single zero was to be found
              where_single_zero = groups_with_single_zero * where_zeros
              # further optimization to avoid computing ProdWithoutZeros
              # if the incoming gradient is 0
              where_gz_not_zero = neq(gz, 0.0)
              # only take ProdWithoutZeros for the groups with single zeros
              # with non-null incoming gradient
              where_to_take_prod_without_zeros = (
                  groups_with_single_zero * where_gz_not_zero
              )
              # preprocess the original input so that we set 0 everywhere
              # except for groups that contain a single zero, to avoid computing
              # multiplications on other groups
              prod_without_zeros_in = where_to_take_prod_without_zeros * prod_in
              # TODO: put lazy switch here, if it'd work
              # this is pretty efficient already (no multiplication if 0), but
              # it'd be even better if we had a lazy if per element
              prod_without_zeros = ProdWithoutZeros(axis=self.axis)(prod_without_zeros_in)
              prod_without_zeros = prod_without_zeros.dimshuffle(new_dims)
  
              groups_without_zeros = eq(sum_where_zeros, 0).dimshuffle(new_dims)
  
              final_grad = switch(
                  groups_without_zeros,
                  grad_case_without_zeros,
                  switch(where_single_zero, prod_without_zeros, 0.0) * gz,
              )
  
              return [final_grad]
  
      def c_code_cache_version(self):
          return (1,)
  
  
  def prod(
      input,
      axis=None,
      dtype=None,
      keepdims=False,
      acc_dtype=None,
      no_zeros_in_input=False,
  ):
      """
      Computes the product along the given axis(es) of a tensor `input`.
  
      When axis is None (the default value), the product is performed
      over the flattened tensor.
  
      For full documentation see ``tensor.elemwise.Prod``.
  
      Parameters
      ----------
      keepdims: bool
          If this is set to True, the axes which are reduced are left in
          the result as dimensions with size one. With this option, the result
          will broadcast correctly against the original tensor.
  
      """
  
      out = Prod(
          axis, dtype=dtype, acc_dtype=acc_dtype, no_zeros_in_input=no_zeros_in_input
      )(input)
  
      if keepdims:
          out = makeKeepDims(input, out, axis)
      return out
  
  
  class MulWithoutZeros(BinaryScalarOp):
      # "identity" here is zero, as in Reduce we don't want to start
      # with reducing (1, something_else): this leads to the erroneous
      # case where a vector of zeros is reduced by binary reductions
      # of (1, 0), which always ends up as 1 (ie. the result for
      # the c version, for the product of [0,0,0], is 1.0)
  
      identity = 0.0
      commutative = True
      associative = True
  
      def impl(self, x, y):
          if x == 0:
              return y
          if y == 0:
              return x
          return x * y
  
      def c_code(self, node, name, inp, out, sub):
          x, y = inp
          (z,) = out
          return (
              "%(z)s = ((%(x)s == 0) ? (%(y)s) : "
              + "((%(y)s == 0) ? (%(x)s) : ((%(y)s)*(%(x)s))) );"
          ) % locals()
  
      def c_code_cache_version(self):
          return (1,)
  
  
  mul_without_zeros = MulWithoutZeros(aes.upcast_out, name="mul_without_zeros")
  
  
  class ProdWithoutZeros(CAReduceDtype):
  
      __props__ = ("axis", "dtype", "acc_dtype")
  
      def __init__(self, axis=None, dtype=None, acc_dtype=None):
          super().__init__(mul_without_zeros, axis=axis, dtype=dtype, acc_dtype=acc_dtype)
  
      def grad(self, inp, grads):
          from aesara.gradient import grad_not_implemented
  
          (a,) = inp
          a_grad = grad_not_implemented(
              self,
              0,
              a,
              "2nd derivatives of `product(a)` is not currently supported."
              "If `a` is guaranteed to contains no zeros, use "
              "`product(a, no_zeros_in_input=True)`.",
          )
          return [a_grad]
  
  
  def any(x, axis=None, keepdims=False):
      out = Any(axis)(x)
  
      if keepdims:
          out = makeKeepDims(x, out, axis)
      return out
  
  
  def all(x, axis=None, keepdims=False):
      out = All(axis)(x)
  
      if keepdims:
          out = makeKeepDims(x, out, axis)
      return out
  
  
  def ptp(a, axis=None):
      """
      Range of values (maximum - minimum) along an axis.
  
      The name of the function comes from the acronym for peak to peak.
  
      Parameters
      ----------
      a
          Input tensor.
      axis
          Axis along which to find the peaks. By default, flatten the array.
  
      Returns
      -------
      array
          A new array holding the result.
  
      """
  
      a = as_tensor_variable(a)
  
      out = max(a, axis) - min(a, axis)
  
      return out
  
  
  def power(x, y):
      return x ** y
  
  
  def logaddexp(*xs):
      """Logarithm of the sum of exponentiations of the inputs.
  
      See ``numpy.logaddexp``.
  
      Parameters
      ----------
      xs : symbolic tensors
          Input
  
      Returns
      -------
      tensor
  
      """
  
      return log(add(*[exp(x) for x in xs]))
  
  
  def logsumexp(x, axis=None, keepdims=False):
      """Compute the log of the sum of exponentials of input elements.
  
      See ``scipy.special.logsumexp``.
  
      Parameters
      ----------
      x : symbolic tensor
          Input
  
      axis : None or int or tuple of ints, optional
          Axis or axes over which the sum is taken. By default axis is None,
          and all elements are summed.
  
      keepdims : bool, optional
          If this is set to True, the axes which are reduced are left in the
          result as dimensions with size one. With this option, the result will
          broadcast correctly against the original array.
  
      Returns
      -------
      tensor
  
      """
  
      return log(sum(exp(x), axis=axis, keepdims=keepdims))
  
  
  __all__ = [
      "max_and_argmax",
      "max",
      "argmax",
      "min",
      "argmin",
      "smallest",
      "largest",
      "lt",
      "gt",
      "le",
      "ge",
      "eq",
      "neq",
      "isnan",
      "isinf",
      "allclose",
      "isclose",
      "and_",
      "bitwise_and",
      "or_",
      "bitwise_or",
      "xor",
      "bitwise_xor",
      "invert",
      "bitwise_not",
      "abs",
      "abs_",
      "exp",
      "exp2",
      "expm1",
      "neg",
      "reciprocal",
      "inv",
      "log",
      "log2",
      "log10",
      "log1p",
      "sgn",
      "ceil",
      "floor",
      "trunc",
      "iround",
      "round",
      "round_half_to_even",
      "round_half_away_from_zero",
      "sqr",
      "square",
      "cov",
      "sqrt",
      "deg2rad",
      "rad2deg",
      "cos",
      "arccos",
      "sin",
      "arcsin",
      "tan",
      "arctan",
      "arctan2",
      "cosh",
      "arccosh",
      "sinh",
      "arcsinh",
      "tanh",
      "arctanh",
      "erf",
      "erfc",
      "erfcx",
      "erfinv",
      "erfcinv",
      "gamma",
      "gammaln",
      "psi",
      "tri_gamma",
      "chi2sf",
      "gammainc",
      "gammaincc",
      "gammau",
      "gammal",
      "j0",
      "j1",
      "jv",
      "i0",
      "i1",
      "iv",
      "sigmoid",
      "expit",
      "softplus",
      "log1pexp",
      "log1mexp",
      "betainc",
      "real",
      "imag",
      "angle",
      "complex",
      "conj",
      "complex_from_polar",
      "sum",
      "prod",
      "mean",
      "var",
      "std",
      "std",
      "maximum",
      "minimum",
      "divmod",
      "add",
      "sub",
      "mul",
      "true_div",
      "int_div",
      "floor_div",
      "ceil_intdiv",
      "mod",
      "pow",
      "clip",
      "dot",
      "dense_dot",
      "tensordot",
      "outer",
      "any",
      "all",
      "ptp",
      "power",
      "logaddexp",
      "logsumexp",
  ]
  
apis: numpy.asarray numpy.max numpy.min numpy.zeros
apis: numpy.allclose numpy.dot numpy.maximum numpy.asarray numpy.argmax numpy.int64 numpy.iinfo numpy.max numpy.prod numpy.array numpy.concatenate numpy.isscalar numpy.mean