code: def local_logsoftmax(fgraph, node):
      """
      Detect Log(Softmax(x)) and replace it with LogSoftmax(x)
  
      This also lifts Subtensor or Dimshuffle operations that could be in between log and softmax
  
      Note: only forward pass is affected
      """
  
      def find_softmax_under_lifteable_ops(inp_node, ops_to_lift):
          if inp_node is None:
              return
  
          if isinstance(inp_node.op, Softmax):
              return inp_node
  
          if isinstance(inp_node.op, subtensor_ops):
              ops_to_lift.append((inp_node.op, inp_node.inputs[1:]))
              return find_softmax_under_lifteable_ops(
                  inp_node.inputs[0].owner, ops_to_lift
              )
  
          if isinstance(inp_node.op, DimShuffle):
              ops_to_lift.append((inp_node.op, ()))
              return find_softmax_under_lifteable_ops(
                  inp_node.inputs[0].owner, ops_to_lift
              )
  
      ops_to_lift = []
      softmax_node = find_softmax_under_lifteable_ops(
          node.inputs[0].owner, ops_to_lift
      )
  
      if softmax_node is None:
          return
  
      ret = log_softmax(
          softmax_node.inputs[0], axis=softmax_node.op.axis
      )
      ret.tag.values_eq_approx = values_eq_approx_remove_inf
  
      # Lift ops that used to be between log and softmax
      for op_to_lift, parameters in reversed(ops_to_lift):
          ret = op_to_lift(ret, *parameters)
  
      copy_stack_trace(node.outputs, ret)
      return [ret]
  
file_path: pytensor/pytensor/tensor/rewriting/special.py
imports: from pytensor.graph.rewriting.basic import copy_stack_trace, node_rewriter
  from pytensor.tensor.elemwise import DimShuffle
  from pytensor.tensor.math import Sum, exp, log, true_div
  from pytensor.tensor.math import sum
  from pytensor.tensor.rewriting.basic import register_stabilize
  from pytensor.tensor.rewriting.math import local_mul_canonizer
  from pytensor.tensor.special import Softmax, SoftmaxGrad, log_softmax
  from pytensor.tensor.subtensor import AdvancedIncSubtensor, AdvancedSubtensor, AdvancedSubtensor1, Subtensor
  from pytensor.tensor.type import values_eq_approx_remove_inf, values_eq_approx_remove_nan
label_desc: Use image node_id 1 for calling a global function with example usage: local_logsoftmax(fgraph, node) and returns: list, inp_node, find_softmax_under_lifteable_ops, find_softmax_under_lifteable_ops