code: def test_jax_CAReduce():
      a_pt = vector("a")
      a_pt.tag.test_value = np.r_[1, 2, 3].astype(config.floatX)
  
      x = pt_sum(a_pt, axis=None)
      x_fg = FunctionGraph([a_pt], [x])
  
      compare_jax_and_py(x_fg, [np.r_[1, 2, 3].astype(config.floatX)])
  
      a_pt = matrix("a")
      a_pt.tag.test_value = np.c_[[1, 2, 3], [1, 2, 3]].astype(
          config.floatX
      )
  
      x = pt_sum(a_pt, axis=0)
      x_fg = FunctionGraph([a_pt], [x])
  
      compare_jax_and_py(
          x_fg, [np.c_[[1, 2, 3], [1, 2, 3]].astype(config.floatX)]
      )
  
      x = pt_sum(a_pt, axis=1)
      x_fg = FunctionGraph([a_pt], [x])
  
      compare_jax_and_py(
          x_fg, [np.c_[[1, 2, 3], [1, 2, 3]].astype(config.floatX)]
      )
  
      a_pt = matrix("a")
      a_pt.tag.test_value = np.c_[[1, 2, 3], [1, 2, 3]].astype(
          config.floatX
      )
  
      x = prod(a_pt, axis=0)
      x_fg = FunctionGraph([a_pt], [x])
  
      compare_jax_and_py(
          x_fg, [np.c_[[1, 2, 3], [1, 2, 3]].astype(config.floatX)]
      )
  
      x = pt_all(a_pt)
      x_fg = FunctionGraph([a_pt], [x])
  
      compare_jax_and_py(
          x_fg, [np.c_[[1, 2, 3], [1, 2, 3]].astype(config.floatX)]
      )
  
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 3 for calling a global function with example usage: test_jax_CAReduce() without return types