code: def test_jax_Subtensor_boolean_mask_reexpressible():
      """Summing values with boolean indexing.
  
      This test ensures that the sum of an `AdvancedSubtensor` `Op`s with boolean
      indexing is replaced with the sum of an equivalent `Switch` `Op`, using the
      `jax_boolean_indexing_sum` rewrite.
  
      JAX forces users to re-express this logic manually, so this is an
      improvement over its user interface.
  
      """
      x_pt = pt.matrix("x")
      out_pt = x_pt[x_pt < 0].sum()
      out_fg = FunctionGraph([x_pt], [out_pt])
      compare_jax_and_py(
          out_fg, [np.arange(25).reshape(5, 5).astype(config.floatX)]
      )
  
file_path: pytensor/tests/link/jax/test_subtensor.py
imports: import numpy
  import pytest
  import pytensor.tensor
  from pytensor.configdefaults import config
  from pytensor.graph.fg import FunctionGraph
  from pytensor.tensor import subtensor
  from pytensor.tensor.rewriting.jax import boolean_indexing_set_or_inc, boolean_indexing_sum
  from tests.link.jax.test_basic import compare_jax_and_py
label_desc: Use image node_id 4 for calling a global function with example usage: test_jax_Subtensor_boolean_mask_reexpressible() without return types