text: import operator
  import warnings
  from functools import singledispatch
  
  import numba
  import numba.np.unsafe.ndarray as numba_ndarray
  import numpy as np
  import scipy
  import scipy.special
  from llvmlite.llvmpy.core import Type as llvm_Type
  from numba import types
  from numba.core.errors import TypingError
  from numba.extending import box
  
  from aesara.compile.ops import DeepCopyOp
  from aesara.graph.basic import Apply
  from aesara.graph.fg import FunctionGraph
  from aesara.graph.type import Type
  from aesara.ifelse import IfElse
  from aesara.link.utils import (
      compile_function_src,
      fgraph_to_python,
      unique_name_generator,
  )
  from aesara.scalar.basic import Scalar
  from aesara.scalar.math import Softplus
  from aesara.tensor.blas import BatchedDot
  from aesara.tensor.math import Dot
  from aesara.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape
  from aesara.tensor.slinalg import Cholesky, Solve
  from aesara.tensor.subtensor import (
      AdvancedIncSubtensor,
      AdvancedIncSubtensor1,
      AdvancedSubtensor,
      AdvancedSubtensor1,
      IncSubtensor,
      Subtensor,
  )
  from aesara.tensor.type import TensorType
  from aesara.tensor.type_other import MakeSlice
  
  
  def get_numba_type(
      aesara_type: Type, layout: str = "A", force_scalar: bool = False
  ) -> numba.types.Type:
      """Create a Numba type object for a ``Type``."""
  
      if isinstance(aesara_type, TensorType):
          dtype = aesara_type.numpy_dtype
          numba_dtype = numba.from_dtype(dtype)
          if force_scalar:
              return numba_dtype
          return numba.types.Array(numba_dtype, aesara_type.ndim, layout)
      elif isinstance(aesara_type, Scalar):
          dtype = np.dtype(aesara_type.dtype)
          numba_dtype = numba.from_dtype(dtype)
          return numba_dtype
      else:
          raise NotImplementedError(f"Numba type not implemented for {aesara_type}")
  
  
  def create_numba_signature(node: Apply, force_scalar: bool = False) -> numba.types.Type:
      """Create a Numba type for the signature of an ``Apply`` node."""
      input_types = []
      for inp in node.inputs:
          input_types.append(get_numba_type(inp.type, force_scalar=force_scalar))
  
      output_types = []
      for out in node.outputs:
          output_types.append(get_numba_type(out.type, force_scalar=force_scalar))
  
      if len(output_types) > 1:
          return numba.types.Tuple(output_types)(*input_types)
      elif len(output_types) == 1:
          return output_types[0](*input_types)
      else:
          return numba.types.void(*input_types)
  
  
  def slice_new(self, start, stop, step):
      fnty = llvm_Type.function(self.pyobj, [self.pyobj, self.pyobj, self.pyobj])
      fn = self._get_function(fnty, name="PySlice_New")
      return self.builder.call(fn, [start, stop, step])
  
  
  def enable_slice_boxing():
      """Enable boxing for Numba's native ``slice``s.
  
      TODO: this can be removed when https://github.com/numba/numba/pull/6939 is
      merged and a release is made.
      """
  
      @box(types.SliceType)
      def box_slice(typ, val, c):
          """Implement boxing for ``slice`` objects in Numba.
  
          This makes it possible to return an Numba's internal representation of a
          ``slice`` object as a proper ``slice`` to Python.
          """
  
          start = c.box(types.int64, c.builder.extract_value(val, 0))
          stop = c.box(types.int64, c.builder.extract_value(val, 1))
          if typ.has_step:
              step = c.box(types.int64, c.builder.extract_value(val, 2))
          else:
              step = c.pyapi.get_null_object()
  
          slice_val = slice_new(c.pyapi, start, stop, step)
  
          return slice_val
  
      @numba.extending.overload(operator.contains)
      def in_seq_empty_tuple(x, y):
          if isinstance(x, types.Tuple) and not x.types:
              return lambda x, y: False
  
  
  enable_slice_boxing()
  
  
  @numba.generated_jit(nopython=True)
  def to_scalar(x):
      if isinstance(x, (numba.types.Number, numba.types.Boolean)):
          return lambda x: x
      elif isinstance(x, numba.types.Array):
          return lambda x: x.item()
      else:
          raise TypingError(f"{x} must be a scalar compatible type.")
  
  
  def enable_slice_literals():
      """Enable lowering for ``SliceLiteral``s.
  
      TODO: This can be removed once https://github.com/numba/numba/pull/6996 is merged
      and a release is made.
      """
      from numba.core import types
      from numba.core.datamodel.models import SliceModel
      from numba.core.datamodel.registry import register_default
      from numba.core.imputils import lower_cast, lower_constant
      from numba.core.types.misc import SliceLiteral
      from numba.cpython.slicing import get_defaults
  
      register_default(numba.types.misc.SliceLiteral)(SliceModel)
  
      @property
      def key(self):
          return self.name
  
      SliceLiteral.key = key
  
      def make_slice_from_constant(context, builder, ty, pyval):
          sli = context.make_helper(builder, ty)
          lty = context.get_value_type(types.intp)
  
          (
              default_start_pos,
              default_start_neg,
              default_stop_pos,
              default_stop_neg,
              default_step,
          ) = [context.get_constant(types.intp, x) for x in get_defaults(context)]
  
          step = pyval.step
          if step is None:
              step_is_neg = False
              step = default_step
          else:
              step_is_neg = step < 0
              step = lty(step)
  
          start = pyval.start
          if start is None:
              if step_is_neg:
                  start = default_start_neg
              else:
                  start = default_start_pos
          else:
              start = lty(start)
  
          stop = pyval.stop
          if stop is None:
              if step_is_neg:
                  stop = default_stop_neg
              else:
                  stop = default_stop_pos
          else:
              stop = lty(stop)
  
          sli.start = start
          sli.stop = stop
          sli.step = step
  
          return sli._getvalue()
  
      @lower_constant(numba.types.SliceType)
      def constant_slice(context, builder, ty, pyval):
          if isinstance(ty, types.Literal):
              typ = ty.literal_type
          else:
              typ = ty
  
          return make_slice_from_constant(context, builder, typ, pyval)
  
      @lower_cast(numba.types.misc.SliceLiteral, numba.types.SliceType)
      def cast_from_literal(context, builder, fromty, toty, val):
          return make_slice_from_constant(
              context,
              builder,
              toty,
              fromty.literal_value,
          )
  
  
  enable_slice_literals()
  
  
  def create_tuple_creator(f, n):
      """Construct a compile-time ``tuple``-comprehension-like loop.
  
      See https://github.com/numba/numba/issues/2771#issuecomment-414358902
      """
      assert n > 0
  
      f = numba.njit(f)
  
      @numba.njit
      def creator(args):
          return (f(0, *args),)
  
      for i in range(1, n):
  
          @numba.njit
          def creator(args, creator=creator, i=i):
              return creator(args) + (f(i, *args),)
  
      return numba.njit(lambda *args: creator(args))
  
  
  def create_tuple_string(x):
      args = ", ".join(x + ([""] if len(x) == 1 else []))
      return f"({args})"
  
  
  def create_arg_string(x):
      args = ", ".join(x)
      return args
  
  
  @singledispatch
  def numba_typify(data, dtype=None, **kwargs):
      return data
  
  
  @singledispatch
  def numba_funcify(op, node=None, storage_map=None, **kwargs):
      """Create a Numba compatible function from an Aesara `Op`."""
  
      warnings.warn(
          f"Numba will use object mode to run {op}'s perform method",
          UserWarning,
      )
  
      n_outputs = len(node.outputs)
  
      if n_outputs > 1:
          ret_sig = numba.types.Tuple([get_numba_type(o.type) for o in node.outputs])
      else:
          ret_sig = get_numba_type(node.outputs[0].type)
  
      @numba.njit
      def perform(*inputs):
          with numba.objmode(ret=ret_sig):
              outputs = [[None] for i in range(n_outputs)]
              op.perform(node, inputs, outputs)
              outputs = tuple([o[0] for o in outputs])
              if n_outputs == 1:
                  ret = outputs[0]
              else:
                  ret = outputs
          return ret
  
      return perform
  
  
  @numba_funcify.register(FunctionGraph)
  def numba_funcify_FunctionGraph(
      fgraph,
      node=None,
      fgraph_name="numba_funcified_fgraph",
      **kwargs,
  ):
      return fgraph_to_python(
          fgraph,
          numba_funcify,
          type_conversion_fn=numba_typify,
          fgraph_name=fgraph_name,
          **kwargs,
      )
  
  
  def create_index_func(node, objmode=False):
      """Create a Python function that assembles and uses an index on an array."""
  
      def convert_indices(indices, entry):
          if indices and isinstance(entry, Type):
              rval = indices.pop(0)
              return rval.auto_name
          elif isinstance(entry, slice):
              return (
                  f"slice({convert_indices(indices, entry.start)}, "
                  f"{convert_indices(indices, entry.stop)}, "
                  f"{convert_indices(indices, entry.step)})"
              )
          elif isinstance(entry, type(None)):
              return "None"
          else:
              raise ValueError()
  
      set_or_inc = isinstance(
          node.op, (IncSubtensor, AdvancedIncSubtensor1, AdvancedIncSubtensor)
      )
      index_start_idx = 1 + int(set_or_inc)
  
      unique_names = unique_name_generator(
          ["subtensor", "incsubtensor", "z"], suffix_sep="_"
      )
  
      input_names = [unique_names(v, force_unique=True) for v in node.inputs]
      op_indices = list(node.inputs[index_start_idx:])
      idx_list = getattr(node.op, "idx_list", None)
  
      indices_creation_src = (
          tuple(convert_indices(op_indices, idx) for idx in idx_list)
          if idx_list
          else tuple(input_names[index_start_idx:])
      )
  
      if len(indices_creation_src) == 1:
          indices_creation_src = f"indices = ({indices_creation_src[0]},)"
      else:
          indices_creation_src = ", ".join(indices_creation_src)
          indices_creation_src = f"indices = ({indices_creation_src})"
  
      if set_or_inc:
          fn_name = "incsubtensor"
          if node.op.inplace:
              index_prologue = f"z = {input_names[0]}"
          else:
              index_prologue = f"z = np.copy({input_names[0]})"
  
          if node.inputs[1].ndim == 0:
              # TODO FIXME: This is a hack to get around a weird Numba typing
              # issue.  See https://github.com/numba/numba/issues/6000
              y_name = f"{input_names[1]}.item()"
          else:
              y_name = input_names[1]
  
          if node.op.set_instead_of_inc:
              index_body = f"z[indices] = {y_name}"
          else:
              index_body = f"z[indices] += {y_name}"
      else:
          fn_name = "subtensor"
          index_prologue = ""
          index_body = f"z = {input_names[0]}[indices]"
  
      if objmode:
          output_var = node.outputs[0]
  
          if not set_or_inc:
              # Since `z` is being "created" while in object mode, it's
              # considered an "outgoing" variable and needs to be manually typed
              output_sig = f"z='{output_var.dtype}[{', '.join([':'] * output_var.ndim)}]'"
          else:
              output_sig = ""
  
          index_body = f"""
      with objmode({output_sig}):
          {index_body}
          """
  
      subtensor_def_src = f"""
  def {fn_name}({", ".join(input_names)}):
      {index_prologue}
      {indices_creation_src}
      {index_body}
      return z
      """
  
      return subtensor_def_src
  
  
  @numba_funcify.register(Subtensor)
  @numba_funcify.register(AdvancedSubtensor)
  @numba_funcify.register(AdvancedSubtensor1)
  def numba_funcify_Subtensor(op, node, **kwargs):
  
      subtensor_def_src = create_index_func(
          node, objmode=isinstance(op, AdvancedSubtensor)
      )
  
      global_env = {"np": np, "objmode": numba.objmode}
  
      subtensor_fn = compile_function_src(subtensor_def_src, "subtensor", global_env)
  
      return numba.njit(subtensor_fn)
  
  
  @numba_funcify.register(IncSubtensor)
  @numba_funcify.register(AdvancedIncSubtensor)
  @numba_funcify.register(AdvancedIncSubtensor1)
  def numba_funcify_IncSubtensor(op, node, **kwargs):
  
      incsubtensor_def_src = create_index_func(
          node, objmode=isinstance(op, AdvancedIncSubtensor)
      )
  
      global_env = {"np": np, "objmode": numba.objmode}
  
      incsubtensor_fn = compile_function_src(
          incsubtensor_def_src, "incsubtensor", global_env
      )
  
      return numba.njit(incsubtensor_fn)
  
  
  @numba_funcify.register(DeepCopyOp)
  def numba_funcify_DeepCopyOp(op, node, **kwargs):
  
      # Scalars are apparently returned as actual Python scalar types and not
      # NumPy scalars, so we need two separate Numba functions for each case.
  
      # The type can also be RandomType with no ndims
      if not hasattr(node.outputs[0].type, "ndim") or node.outputs[0].type.ndim == 0:
          # TODO: Do we really need to compile a pass-through function like this?
          @numba.njit(inline="always")
          def deepcopyop(x):
              return x
  
      else:
  
          @numba.njit(inline="always")
          def deepcopyop(x):
              return x.copy()
  
      return deepcopyop
  
  
  @numba_funcify.register(MakeSlice)
  def numba_funcify_MakeSlice(op, **kwargs):
      @numba.njit
      def makeslice(*x):
          return slice(*x)
  
      return makeslice
  
  
  @numba_funcify.register(Shape)
  def numba_funcify_Shape(op, **kwargs):
      @numba.njit(inline="always")
      def shape(x):
          return np.asarray(np.shape(x))
  
      return shape
  
  
  @numba_funcify.register(Shape_i)
  def numba_funcify_Shape_i(op, **kwargs):
      i = op.i
  
      @numba.njit(inline="always")
      def shape_i(x):
          return np.shape(x)[i]
  
      return shape_i
  
  
  @numba.extending.intrinsic
  def direct_cast(typingctx, val, typ):
  
      if isinstance(typ, numba.types.TypeRef):
          casted = typ.instance_type
      elif isinstance(typ, numba.types.DTypeSpec):
          casted = typ.dtype
      else:
          casted = typ
  
      sig = casted(casted, typ)
  
      def codegen(context, builder, signature, args):
          val, _ = args
          context.nrt.incref(builder, signature.return_type, val)
          return val
  
      return sig, codegen
  
  
  @numba_funcify.register(Reshape)
  def numba_funcify_Reshape(op, **kwargs):
      ndim = op.ndim
  
      if ndim == 0:
  
          @numba.njit(inline="always")
          def reshape(x, shape):
              return x.item()
  
      else:
  
          @numba.njit(inline="always")
          def reshape(x, shape):
              # TODO: Use this until https://github.com/numba/numba/issues/7353 is closed.
              return np.reshape(
                  np.ascontiguousarray(np.asarray(x)),
                  numba_ndarray.to_fixed_tuple(shape, ndim),
              )
  
      return reshape
  
  
  @numba_funcify.register(SpecifyShape)
  def numba_funcify_SpecifyShape(op, **kwargs):
      @numba.njit
      def specifyshape(x, shape):
          assert np.array_equal(x.shape, shape)
          return x
  
      return specifyshape
  
  
  def int_to_float_fn(inputs, out_dtype):
      """Create a Numba function that converts integer and boolean ``ndarray``s to floats."""
  
      if any(i.type.numpy_dtype.kind in "ib" for i in inputs):
  
          args_dtype = np.dtype(f"f{out_dtype.itemsize}")
  
          @numba.njit(inline="always")
          def inputs_cast(x):
              return x.astype(args_dtype)
  
      else:
          args_dtype_sz = max([_arg.type.numpy_dtype.itemsize for _arg in inputs])
          args_dtype = np.dtype(f"f{args_dtype_sz}")
  
          @numba.njit(inline="always")
          def inputs_cast(x):
              return x.astype(args_dtype)
  
      return inputs_cast
  
  
  @numba_funcify.register(Dot)
  def numba_funcify_Dot(op, node, **kwargs):
      # Numba's `np.dot` does not support integer dtypes, so we need to cast to
      # float.
  
      out_dtype = node.outputs[0].type.numpy_dtype
      inputs_cast = int_to_float_fn(node.inputs, out_dtype)
  
      @numba.njit(inline="always")
      def dot(x, y):
          return np.asarray(np.dot(inputs_cast(x), inputs_cast(y))).astype(out_dtype)
  
      return dot
  
  
  @numba_funcify.register(Softplus)
  def numba_funcify_Softplus(op, node, **kwargs):
  
      x_dtype = np.dtype(node.inputs[0].dtype)
  
      @numba.njit
      def softplus(x):
          if x < -37.0:
              return direct_cast(np.exp(x), x_dtype)
          elif x < 18.0:
              return direct_cast(np.log1p(np.exp(x)), x_dtype)
          elif x < 33.3:
              return direct_cast(x + np.exp(-x), x_dtype)
          else:
              return direct_cast(x, x_dtype)
  
      return softplus
  
  
  @numba_funcify.register(Cholesky)
  def numba_funcify_Cholesky(op, node, **kwargs):
      lower = op.lower
  
      out_dtype = node.outputs[0].type.numpy_dtype
  
      if lower:
  
          inputs_cast = int_to_float_fn(node.inputs, out_dtype)
  
          @numba.njit(inline="always")
          def cholesky(a):
              return np.linalg.cholesky(inputs_cast(a)).astype(out_dtype)
  
      else:
          # TODO: Use SciPy's BLAS/LAPACK Cython wrappers.
  
          warnings.warn(
              (
                  "Numba will use object mode to allow the "
                  "`lower` argument to `scipy.linalg.cholesky`."
              ),
              UserWarning,
          )
  
          ret_sig = get_numba_type(node.outputs[0].type)
  
          @numba.njit
          def cholesky(a):
              with numba.objmode(ret=ret_sig):
                  ret = scipy.linalg.cholesky(a, lower=lower).astype(out_dtype)
              return ret
  
      return cholesky
  
  
  @numba_funcify.register(Solve)
  def numba_funcify_Solve(op, node, **kwargs):
  
      assume_a = op.assume_a
      # check_finite = op.check_finite
  
      if assume_a != "gen":
  
          lower = op.lower
  
          warnings.warn(
              (
                  "Numba will use object mode to allow the "
                  "`compute_uv` argument to `numpy.linalg.svd`."
              ),
              UserWarning,
          )
  
          ret_sig = get_numba_type(node.outputs[0].type)
  
          @numba.njit
          def solve(a, b):
              with numba.objmode(ret=ret_sig):
                  ret = scipy.linalg.solve_triangular(
                      a,
                      b,
                      lower=lower,
                      # check_finite=check_finite
                  )
              return ret
  
      else:
          out_dtype = node.outputs[0].type.numpy_dtype
          inputs_cast = int_to_float_fn(node.inputs, out_dtype)
  
          @numba.njit(inline="always")
          def solve(a, b):
              return np.linalg.solve(
                  inputs_cast(a),
                  inputs_cast(b),
                  # assume_a=assume_a,
                  # check_finite=check_finite,
              ).astype(out_dtype)
  
      return solve
  
  
  @numba_funcify.register(BatchedDot)
  def numba_funcify_BatchedDot(op, node, **kwargs):
      dtype = node.outputs[0].type.numpy_dtype
  
      @numba.njit
      def batched_dot(x, y):
          shape = x.shape[:-1] + y.shape[2:]
          z0 = np.empty(shape, dtype=dtype)
          for i in range(z0.shape[0]):
              z0[i] = np.dot(x[i], y[i])
  
          return z0
  
      return batched_dot
  
  
  # NOTE: The remaining `aesara.tensor.blas` `Op`s appear unnecessary, because
  # they're only used to optimize basic `Dot` nodes, and those GEMV and GEMM
  # optimizations are apparently already performed by Numba
  
  
  @numba_funcify.register(IfElse)
  def numba_funcify_IfElse(op, **kwargs):
      n_outs = op.n_outs
  
      if n_outs > 1:
  
          @numba.njit
          def ifelse(cond, *args):
              if cond:
                  res = args[:n_outs]
              else:
                  res = args[n_outs:]
  
              return res
  
      else:
  
          @numba.njit
          def ifelse(cond, *args):
              if cond:
                  res = args[:n_outs]
              else:
                  res = args[n_outs:]
  
              return res[0]
  
      return ifelse
  