docstr: Test the `Logsoftmax`'s grad substitution.
  
  Check that ``Log(Softmax(x))``'s grad is substituted with ``Logsoftmax(x)``'s
  grad and that the new operation does not explode for big inputs.
  Note that only the grad is checked.
docstr_tok: ["Test","the","`","Logsoftmax","`","'s","grad","substitution",".","Check","that","``","Log","(","Softmax","(","x",")",")","``","'s","grad","is","substituted","with","``","Logsoftmax","(","x",")","``","'s","grad","and","that","the","new","operation","does","not","explode","for","big","inputs",".","Note","that","only","the","grad","is","checked","."]
code: def test_local_logsoftmax_grad_rewrite(self, axis):
      """Test the `Logsoftmax`'s grad substitution.
  
      Check that ``Log(Softmax(x))``'s grad is substituted with ``Logsoftmax(x)``'s
      grad and that the new operation does not explode for big inputs.
      Note that only the grad is checked.
      """
  
      m = config.mode
      m = get_mode(m).including("stabilize")
      m.check_isfinite = False
      # some inputs that are large to make the gradient explode in the non
      # rewritten case
      rng = np.random.default_rng(utt.fetch_seed())
      a = np.exp(10 * rng.random((5, 10)).astype(config.floatX))
  
      def myfunc(x):
          sm = softmax(x, axis=axis)
          logsm = log(sm)
          return logsm
  
      # We set step to 0.1 because for big values we need a big epsilon
      utt.verify_grad(myfunc, [a], eps=0.1, mode=m)
      sa = shared(a)
      f = FunctionGraph([sa], [myfunc(sa)])
      _fast_run_rewrites(f)
      assert check_stack_trace(f, ops_to_check="all")
  
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 3 for calling the TestLogSoftmaxRewrites obj's underlying member method code with example usage: obj.test_local_logsoftmax_grad_rewrite(axis) and returns: logsm