docstr: Test that stabilization is introduced even when subtensor or dimshuffle operations
  are present between log and softmax.
docstr_tok: ["Test","that","stabilization","is","introduced","even","when","subtensor","or","dimshuffle","operations","are","present","between","log","and","softmax","."]
code: def test_logsoftmax_subtensor_dimshuffle(self, axis, idx0, idx1):
      """Test that stabilization is introduced even when subtensor or dimshuffle operations
      are present between log and softmax.
      """
      logit_p = matrix("logit_p")
      p = softmax(logit_p, axis=axis)
      p_indexed = p[(idx0, idx1)]
      out = log(p_indexed)
  
      # Don't waste time with C compilation
      with config.change_flags(cxx=""):
          mode = get_mode(None).including("stabilize")
          fn = pytensor.function([logit_p], out, mode=mode)
  
      assert not any(
          isinstance(node.op, Softmax)
          for node in fn.maker.fgraph.apply_nodes
      )
  
      # This range would lead to underflow to -inf without the stabilization
      test_logit_p = np.array(
          [[-10.0, -10.0, 999.0], [999.0, 990.0, -10.0]],
          dtype=config.floatX,
      )
      np.testing.assert_allclose(
          fn(logit_p=test_logit_p),
          scipy.special.log_softmax(test_logit_p, axis=axis)[
              (idx0, idx1)
          ],
      )
  
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 2 for calling the TestLogSoftmaxRewrites obj's underlying member method code with example usage: obj.test_logsoftmax_subtensor_dimshuffle(axis, idx0, idx1) without return types