code: def test_jax_IncSubtensor():
      rng = np.random.default_rng(213234)
  
      x_np = rng.uniform(-1, 1, size=(3, 4, 5)).astype(config.floatX)
      x_pt = pt.constant(
          np.arange(3 * 4 * 5).reshape((3, 4, 5)).astype(config.floatX)
      )
  
      # "Set" basic indices
      st_pt = pt.as_tensor_variable(
          np.array(-10.0, dtype=config.floatX)
      )
      out_pt = pt_subtensor.set_subtensor(x_pt[1, 2, 3], st_pt)
      assert isinstance(out_pt.owner.op, pt_subtensor.IncSubtensor)
      out_fg = FunctionGraph([], [out_pt])
      compare_jax_and_py(out_fg, [])
  
      st_pt = pt.as_tensor_variable(
          np.r_[-1.0, 0.0].astype(config.floatX)
      )
      out_pt = pt_subtensor.set_subtensor(x_pt[:2, 0, 0], st_pt)
      assert isinstance(out_pt.owner.op, pt_subtensor.IncSubtensor)
      out_fg = FunctionGraph([], [out_pt])
      compare_jax_and_py(out_fg, [])
  
      out_pt = pt_subtensor.set_subtensor(x_pt[0, 1:3, 0], st_pt)
      assert isinstance(out_pt.owner.op, pt_subtensor.IncSubtensor)
      out_fg = FunctionGraph([], [out_pt])
      compare_jax_and_py(out_fg, [])
  
      # "Set" advanced indices
      st_pt = pt.as_tensor_variable(
          rng.uniform(-1, 1, size=(2, 4, 5)).astype(config.floatX)
      )
      out_pt = pt_subtensor.set_subtensor(x_pt[np.r_[0, 2]], st_pt)
      assert isinstance(
          out_pt.owner.op, pt_subtensor.AdvancedIncSubtensor
      )
      out_fg = FunctionGraph([], [out_pt])
      compare_jax_and_py(out_fg, [])
  
      st_pt = pt.as_tensor_variable(
          np.r_[-1.0, 0.0].astype(config.floatX)
      )
      out_pt = pt_subtensor.set_subtensor(x_pt[[0, 2], 0, 0], st_pt)
      assert isinstance(
          out_pt.owner.op, pt_subtensor.AdvancedIncSubtensor
      )
      out_fg = FunctionGraph([], [out_pt])
      compare_jax_and_py(out_fg, [])
  
      # "Set" boolean indices
      mask_pt = pt.constant(x_np > 0)
      out_pt = pt_subtensor.set_subtensor(x_pt[mask_pt], 0.0)
      assert isinstance(
          out_pt.owner.op, pt_subtensor.AdvancedIncSubtensor
      )
      out_fg = FunctionGraph([], [out_pt])
      compare_jax_and_py(out_fg, [])
  
      # "Increment" basic indices
      st_pt = pt.as_tensor_variable(
          np.array(-10.0, dtype=config.floatX)
      )
      out_pt = pt_subtensor.inc_subtensor(x_pt[1, 2, 3], st_pt)
      assert isinstance(out_pt.owner.op, pt_subtensor.IncSubtensor)
      out_fg = FunctionGraph([], [out_pt])
      compare_jax_and_py(out_fg, [])
  
      st_pt = pt.as_tensor_variable(
          np.r_[-1.0, 0.0].astype(config.floatX)
      )
      out_pt = pt_subtensor.inc_subtensor(x_pt[:2, 0, 0], st_pt)
      assert isinstance(out_pt.owner.op, pt_subtensor.IncSubtensor)
      out_fg = FunctionGraph([], [out_pt])
      compare_jax_and_py(out_fg, [])
  
      out_pt = pt_subtensor.set_subtensor(x_pt[0, 1:3, 0], st_pt)
      assert isinstance(out_pt.owner.op, pt_subtensor.IncSubtensor)
      out_fg = FunctionGraph([], [out_pt])
      compare_jax_and_py(out_fg, [])
  
      # "Increment" advanced indices
      st_pt = pt.as_tensor_variable(
          rng.uniform(-1, 1, size=(2, 4, 5)).astype(config.floatX)
      )
      out_pt = pt_subtensor.inc_subtensor(x_pt[np.r_[0, 2]], st_pt)
      assert isinstance(
          out_pt.owner.op, pt_subtensor.AdvancedIncSubtensor
      )
      out_fg = FunctionGraph([], [out_pt])
      compare_jax_and_py(out_fg, [])
  
      st_pt = pt.as_tensor_variable(
          np.r_[-1.0, 0.0].astype(config.floatX)
      )
      out_pt = pt_subtensor.inc_subtensor(x_pt[[0, 2], 0, 0], st_pt)
      assert isinstance(
          out_pt.owner.op, pt_subtensor.AdvancedIncSubtensor
      )
      out_fg = FunctionGraph([], [out_pt])
      compare_jax_and_py(out_fg, [])
  
      # "Increment" boolean indices
      mask_pt = pt.constant(x_np > 0)
      out_pt = pt_subtensor.set_subtensor(x_pt[mask_pt], 1.0)
      assert isinstance(
          out_pt.owner.op, pt_subtensor.AdvancedIncSubtensor
      )
      out_fg = FunctionGraph([], [out_pt])
      compare_jax_and_py(out_fg, [])
  
      st_pt = pt.as_tensor_variable(x_np[[0, 2], 0, :3])
      out_pt = pt_subtensor.set_subtensor(x_pt[[0, 2], 0, :3], st_pt)
      assert isinstance(
          out_pt.owner.op, pt_subtensor.AdvancedIncSubtensor
      )
      out_fg = FunctionGraph([], [out_pt])
      compare_jax_and_py(out_fg, [])
  
      st_pt = pt.as_tensor_variable(x_np[[0, 2], 0, :3])
      out_pt = pt_subtensor.inc_subtensor(x_pt[[0, 2], 0, :3], st_pt)
      assert isinstance(
          out_pt.owner.op, pt_subtensor.AdvancedIncSubtensor
      )
      out_fg = FunctionGraph([], [out_pt])
      compare_jax_and_py(out_fg, [])
  
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 6 for calling a global function with example usage: test_jax_IncSubtensor() without return types