code: def boolean_indexing_sum(fgraph, node):
      """Replace the sum of `AdvancedSubtensor` with exclusively boolean indexing.
  
      JAX cannot JIT-compile functions that use boolean indexing, but can compile
      those expressions that can be re-expressed using `jax.numpy.where`. This
      rewrite re-rexpressed the model on the behalf of the user and thus allows to
      improve upon JAX's API.
  
      """
      operand = node.inputs[0]
  
      if not isinstance(operand, TensorVariable):
          return
  
      # If it's not a scalar reduction, it couldn't have been a pure boolean mask
      if node.outputs[0].ndim != 0:
          return
  
      if operand.owner is None:
          return
  
      if not isinstance(operand.owner.op, AdvancedSubtensor):
          return
  
      # Get out if AdvancedSubtensor has more than a single indexing operation
      if len(operand.owner.inputs) > 2:
          return
  
      [x, cond] = operand.owner.inputs
  
      if not isinstance(cond, TensorVariable):
          return
  
      if not cond.type.dtype == "bool":
          return
  
      # Output must be a scalar, since pure boolean indexing returns a vector
      # No need to worry about axis
      out = pt.sum(pt.where(cond, x, 0))
      return out.owner.outputs
  
file_path: pytensor/pytensor/tensor/rewriting/jax.py
imports: import pytensor.tensor
  from pytensor.compile import optdb
  from pytensor.graph.rewriting.basic import in2out, node_rewriter
  from pytensor.tensor.basic import MakeVector
  from pytensor.tensor.elemwise import DimShuffle
  from pytensor.tensor.math import Sum
  from pytensor.tensor.shape import Reshape
  from pytensor.tensor.subtensor import AdvancedIncSubtensor, AdvancedSubtensor
  from pytensor.tensor.variable import TensorVariable
label_desc: Use image node_id 2 for calling a global function with example usage: boolean_indexing_sum(fgraph, node) and returns: out