code: def local_logsoftmax_grad(fgraph, node):
      """
      Detect Log(Softmax(x))'s grad and replace it with LogSoftmax(x)'s grad
  
      Note: only grad is affected
      """
      if (
          node.inputs[0].owner is not None
          and node.inputs[0].owner.op == true_div
          and len(node.inputs[0].owner.inputs) >= 2
          and node.inputs[0].owner.inputs[1].owner is not None
          and isinstance(
              node.inputs[0].owner.inputs[1].owner.op, Softmax
          )
          and node.inputs[1] == node.inputs[0].owner.inputs[1]
          and not (
              # skip if it will be optimized by
              # local_advanced_indexing_crossentropy_onehot_grad
              node.inputs[0].owner.op == true_div
              and node.inputs[0].owner.inputs[0].owner is not None
              and isinstance(
                  node.inputs[0].owner.inputs[0].owner.op,
                  AdvancedIncSubtensor,
              )
              # the rewrite only applies to legacy SoftmaxGrad
              and node.op == SoftmaxGrad(axis=-1)
              and node.inputs[0].owner.inputs[1].ndim == 2
          )
      ):
          # get parameters from unoptimized op
          grads, sm = node.inputs[0].owner.inputs
          ret = (
              grads
              - pt_sum(grads, axis=sm.owner.op.axis, keepdims=True) * sm
          )
          ret.tag.values_eq_approx = values_eq_approx_remove_nan
          copy_stack_trace(node.outputs[0], 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 2 for calling a global function with example usage: local_logsoftmax_grad(fgraph, node) and returns: list