code: def test_logsumexp_benchmark(size, axis, benchmark):
      X = pt.matrix("X")
      X_max = pt.max(X, axis=axis, keepdims=True)
      X_max = pt.switch(pt.isinf(X_max), 0, X_max)
      X_lse = (
          pt.log(pt.sum(pt.exp(X - X_max), axis=axis, keepdims=True))
          + X_max
      )
  
      rng = np.random.default_rng(23920)
      X_val = rng.normal(size=size)
  
      X_lse_fn = pytensor.function([X], X_lse, mode="JAX")
  
      # JIT compile first
      _ = X_lse_fn(X_val)
  
      res = benchmark(X_lse_fn, X_val)
  
      exp_res = scipy.special.logsumexp(X_val, axis=axis, keepdims=True)
      np.testing.assert_array_almost_equal(res, exp_res)
  
imports: import numpy
  import pytest
  import scipy.special
  import pytensor
  import pytensor.tensor
  from pytensor.compile import get_mode
  from pytensor.configdefaults import config
  from pytensor.graph.fg import FunctionGraph
  from pytensor.graph.op import get_test_value
  from pytensor.tensor import elemwise
  from pytensor.tensor.math import all
  from pytensor.tensor.math import prod
  from pytensor.tensor.math import sum
  from pytensor.tensor.special import SoftmaxGrad, log_softmax, softmax
  from pytensor.tensor.type import matrix, tensor, vector, vectors
  from tests.link.jax.test_basic import compare_jax_and_py
  from tests.tensor.test_elemwise import TestElemwise
label_desc: Use image node_id 7 for calling a global function with example usage: test_logsumexp_benchmark(size, axis, benchmark) without return types