docstr: Test the `Logsoftmax` substitution.
  
  Check that ``Log(Softmax(x))`` is substituted with ``Logsoftmax(x)``. Note that
  only the forward pass is checked (i.e., doesn't check the gradient)
docstr_tok: ["Test","the","`","Logsoftmax","`","substitution",".","Check","that","``","Log","(","Softmax","(","x",")",")","``","is","substituted","with","``","Logsoftmax","(","x",")","``",".","Note","that","only","the","forward","pass","is","checked","(","i.e.",",","does","n't","check","the","gradient",")"]
code: def test_local_logsoftmax_rewrite(self, axis):
      """Test the `Logsoftmax` substitution.
  
      Check that ``Log(Softmax(x))`` is substituted with ``Logsoftmax(x)``. Note that
      only the forward pass is checked (i.e., doesn't check the gradient)
      """
  
      x = matrix("x")
      sm = softmax(x, axis=axis)
      logsm = log(sm)
      fgraph = FunctionGraph([x], [logsm])
      _fast_run_rewrites.rewrite(fgraph)
      assert isinstance(fgraph.outputs[0].owner.op, LogSoftmax)
      assert check_stack_trace(fgraph, ops_to_check=LogSoftmax)
      assert check_stack_trace(fgraph, 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 1 for calling the TestLogSoftmaxRewrites obj's underlying member method code with example usage: obj.test_local_logsoftmax_rewrite(axis) without return types