code: def boolean_indexing_set_or_inc(fgraph, node):
      """Replace `AdvancedIncSubtensor` when using boolean indexing using `Switch`.
  
      JAX cannot JIT-compile functions that use boolean indexing to set values in
      an array. A workaround is to re-express this logic using `jax.numpy.where`.
      This rewrite allows to improve upon JAX's API.
  
      """
  
      # Get out if there is more than one indexing group
      if len(node.inputs) != 3:
          return None
  
      op = node.op
      [x, y, cond] = node.inputs
  
      # This rewrite only works when `y` is a scalar, so it can broadcast to the shape of x[cond]
      if y.type.ndim > 0:
          return
  
      if not isinstance(cond, TensorVariable):
          return
  
      if not cond.type.dtype == "bool":
          return
  
      if op.set_instead_of_inc:
          out = pt.where(cond, y, x)
          return out.owner.outputs
      else:
          out = pt.where(cond, x + y, x)
          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 1 for calling a global function with example usage: boolean_indexing_set_or_inc(fgraph, node) and returns: None, out, out