docstr: Checks that the gradient of an expression similar to a ``log(softmax)`` but
  with a different elemwise operation than true_div is not rewritten.
docstr_tok: ["Checks","that","the","gradient","of","an","expression","similar","to","a","``","log","(","softmax",")","``","but","with","a","different","elemwise","operation","than","true_div","is","not","rewritten","."]
code: def test_logsoftmax_grad_true_div_elemwise(self):
      """
      Checks that the gradient of an expression similar to a ``log(softmax)`` but
      with a different elemwise operation than true_div is not rewritten.
      """
  
      x = matrix("x")
      y = log(softmax(x, axis=-1))
      g = pytensor.tensor.grad(y.sum(), x)
  
      softmax_grad_node = g.owner
      assert softmax_grad_node.op == SoftmaxGrad(axis=-1)
      true_div_node = softmax_grad_node.inputs[0].owner
      assert true_div_node.op == true_div
  
      # We replace thk elemwise true_div op by an elemwise add.
      new_g = SoftmaxGrad(axis=-1)(
          add(*true_div_node.inputs), softmax_grad_node.inputs[1]
      )
  
      fgraph = FunctionGraph([x], [new_g])
      _fast_run_rewrites.rewrite(fgraph)
  
      assert SoftmaxGrad(axis=-1) in [n.op for n in fgraph.toposort()]
  
imports: import numpy
  import pytest
  import scipy.special
  import pytensor
  from pytensor import shared
  from pytensor.compile import optdb
  from pytensor.compile.mode import get_mode
  from pytensor.configdefaults import config
  from pytensor.graph.fg import FunctionGraph
  from pytensor.graph.rewriting.basic import check_stack_trace
  from pytensor.graph.rewriting.db import RewriteDatabaseQuery
  from pytensor.tensor.math import add, exp, log, true_div
  from pytensor.tensor.special import LogSoftmax, Softmax, SoftmaxGrad, softmax
  from pytensor.tensor.type import matrix
  from tests import unittest_tools
label_desc: Use image node_id 4 for calling the TestLogSoftmaxRewrites obj's underlying member method code with example usage: obj.test_logsoftmax_grad_true_div_elemwise() without return types