From ef37865771f6a9ae2889b769e4c6623f604a3932 Mon Sep 17 00:00:00 2001
From: Masato Shinokawa <shino16@users.noreply.github.com>
Date: Tue, 9 Dec 2025 09:20:15 -0800
Subject: [PATCH 01/15] Change copy_with_setitem to setitem

---
 thunder/clang/__init__.py    |  4 ++--
 thunder/core/prims.py        |  7 +++----
 thunder/core/transforms.py   | 20 --------------------
 thunder/executors/torchex.py | 22 +++++++++-------------
 thunder/torch/__init__.py    |  4 ++--
 5 files changed, 16 insertions(+), 41 deletions(-)

diff --git a/thunder/clang/__init__.py b/thunder/clang/__init__.py
index 333086b906..dd2044f69f 100644
--- a/thunder/clang/__init__.py
+++ b/thunder/clang/__init__.py
@@ -648,10 +648,10 @@ def wrap_tensor(t: TensorLike, dim_length: int) -> TensorLike:
 
 
 @clangop()
-def copy_with_setitem(a: TensorLike, key, value: TensorLike) -> TensorLike:
+def setitem(a: TensorLike, key, value: TensorLike) -> TensorLike:
     # TODO: do more checking here. We used to have a check
     #     lambda: f"{key=} tries to index more dimensions than {a.ndim=}",
-    return prims.copy_with_setitem(a, key, value)
+    return prims.setitem(a, key, value)
 
 
 # NOTE: currently supported indexing:
diff --git a/thunder/core/prims.py b/thunder/core/prims.py
index 90e1f3689b..b31984b9c2 100644
--- a/thunder/core/prims.py
+++ b/thunder/core/prims.py
@@ -266,7 +266,7 @@ class PrimIDs(Enum):
     SCATTER_ADD = auto()
     TAKE = auto()
     TAKE_ALONG_AXIS = auto()
-    COPY_WITH_SETITEM = auto()
+    SETITEM = auto()
     # Linear algebra prims (Mostly experimental)
     MATMUL = auto()
     _GROUPED_MM = auto()  # Used for grouped matmuls
@@ -3615,12 +3615,11 @@ def take_along_axis_meta(a: TensorProxy, /, index: TensorProxy, dim: int) -> Ten
 take_along_axis = make_prim(PrimIDs.TAKE_ALONG_AXIS, "take_along_axis", meta=take_along_axis_meta)
 
 
-def copy_with_setitem_meta(a: TensorProxy, index, value: TensorProxy) -> TensorProxy:
-    # TODO: port checks from clang, currently there  because of the utilities they need
+def setitem_meta(a: TensorProxy, index, value: TensorProxy | Number | NumberProxy) -> TensorProxy:
     return TensorProxy(like=a)
 
 
-copy_with_setitem = make_prim(PrimIDs.COPY_WITH_SETITEM, "copy_with_setitem", meta=copy_with_setitem_meta)
+setitem = make_prim(PrimIDs.SETITEM, "setitem", meta=setitem_meta, tags=(OpTags.DONT_DCE,))
 
 
 def gather_meta(a: TensorProxy, /, index: TensorProxy, dim: int) -> TensorProxy:
diff --git a/thunder/core/transforms.py b/thunder/core/transforms.py
index 11bf85f4bd..5846f0e6c7 100644
--- a/thunder/core/transforms.py
+++ b/thunder/core/transforms.py
@@ -1442,26 +1442,6 @@ def _maximum_grad(a: TensorProxy, b: TensorProxy, /):
 register_grad(pids.SHAPE, prims.shape)
 
 
-def _copy_with_setitem_grad(a: TensorProxy, index, value: Number | TensorProxy):
-    fwd = prims.copy_with_setitem(a, index, value)
-    g = get_grad(fwd)
-
-    a_grad = prims.copy_with_setitem(g, index, 0)
-    put_grad(a, a_grad)
-
-    if isinstance(value, TensorProxy):
-        value_grad = g[index]
-        # NOTE: `value` could be broadcasted.
-        if not utils.same_shape(value_grad.shape, value.shape):
-            value_grad = sum_to(value_grad, value.shape)
-        put_grad(value, value_grad)
-
-    return fwd
-
-
-register_grad(pids.COPY_WITH_SETITEM, _copy_with_setitem_grad)
-
-
 def _log_sigmoid_grad(
     a: TensorProxy,
 ) -> TensorProxy:
diff --git a/thunder/executors/torchex.py b/thunder/executors/torchex.py
index c67aae82ae..14f3c84b87 100644
--- a/thunder/executors/torchex.py
+++ b/thunder/executors/torchex.py
@@ -1482,19 +1482,6 @@ def _take_along_axis_prim_transform(a: TensorProxy, /, index: TensorProxy, dim:
 _register_implementation(ltorch.scatter_add, checker=_always_executable, execution_transform=_scatter_add_transform)
 _register_implementation(ltorch.take_along_dim, take_along_dim, checker=_always_executable)
 
-# out of place setitem helper
-
-
-def _copy_with_setitem_impl(a, key, value):
-    c = a.clone()
-    c[key] = value
-    return c
-
-
-copy_with_setitem_impl = ex.register_operator(
-    "copy_with_setitem_impl", meta=prims.copy_with_setitem_meta, fn=_copy_with_setitem_impl
-)
-_register_implementation(prims.copy_with_setitem, copy_with_setitem_impl, checker=_always_executable)
 
 #
 # Linear algebra operations
@@ -2376,6 +2363,15 @@ def _copy__impl(copy_from, copy_to, grad_enabled):
 _register_implementation(prims.copy_, copy_, checker=_always_executable)
 
 
+def _setitem_impl(a, key, value):
+    a[key] = value
+    return a
+
+
+setitem = ex.register_operator("setitem", tags=(prims.OpTags.DONT_DCE,), like=ltorch.setitem_, fn=_setitem_impl)
+_register_implementation(prims.setitem, setitem, checker=_always_executable)
+
+
 def _shape_impl(t):
     return t.shape
 
diff --git a/thunder/torch/__init__.py b/thunder/torch/__init__.py
index 10bca59273..394f823f8f 100644
--- a/thunder/torch/__init__.py
+++ b/thunder/torch/__init__.py
@@ -1223,12 +1223,12 @@ def flip(a: TensorLike, /, *dims: int) -> TensorLike:
 # fake out of place variant
 @torchsymbol(id="setitem")
 def setitem(inp, idx, val):
-    return clang.copy_with_setitem(inp, idx, val)
+    raise NotImplementedError
 
 
 @torchsymbol(torch.Tensor.__setitem__, id="setitem_", is_method=True, tags=(prims.OpTags.IN_PLACE,))
 def setitem_(inp, idx, val):
-    return _copy_(inp, setitem(inp, idx, val))
+    return clang.setitem(inp, idx, val)
 
 
 @torchsymbol(torch.Tensor.__getitem__, id="torch.Tensor.__getitem__", method_name="getitem")

From 17c02a881a6c0c2b366db7f855aa60ad2fbeff72 Mon Sep 17 00:00:00 2001
From: Masato Shinokawa <shino16@users.noreply.github.com>
Date: Tue, 9 Dec 2025 19:35:53 -0800
Subject: [PATCH 02/15] Empty tensors are not aliases

---
 thunder/__init__.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/thunder/__init__.py b/thunder/__init__.py
index c0483f8c3a..d87f0bb2a2 100644
--- a/thunder/__init__.py
+++ b/thunder/__init__.py
@@ -418,6 +418,9 @@ def _alias_tensor_of_args_kwargs_dict(*args, **kwargs) -> dict[int, list[int]]:
             # subclasses.
             if type(t) is pytorch.Tensor and t.layout is pytorch.strided:
                 data_ptr = t.untyped_storage().data_ptr()
+                if not data_ptr:
+                    # This happens when t.numel() == 0
+                    continue
                 if data_ptr not in data_ptr_to_tensor_group_index:
                     data_ptr_to_tensor_group_index[data_ptr] = len(data_ptr_to_tensor_group_index)
                 tgi = data_ptr_to_tensor_group_index[data_ptr]

From d063dad8b79fbfbec635a6fb9ee5209082e035d1 Mon Sep 17 00:00:00 2001
From: Masato Shinokawa <shino16@users.noreply.github.com>
Date: Tue, 9 Dec 2025 19:36:03 -0800
Subject: [PATCH 03/15] Add backward

---
 thunder/core/transforms.py | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/thunder/core/transforms.py b/thunder/core/transforms.py
index 5846f0e6c7..d1a247cd05 100644
--- a/thunder/core/transforms.py
+++ b/thunder/core/transforms.py
@@ -1442,6 +1442,26 @@ def _maximum_grad(a: TensorProxy, b: TensorProxy, /):
 register_grad(pids.SHAPE, prims.shape)
 
 
+def _setitem_grad(a: TensorProxy, index, value: Number | TensorProxy):
+    fwd = prims.setitem(a, index, value)
+    g = get_grad(fwd)
+
+    a_grad = prims.setitem(prims.clone(g), index, 0)
+    put_grad(a, a_grad)
+
+    if isinstance(value, TensorProxy):
+        value_grad = g[index]
+        # NOTE: `value` could be broadcasted.
+        if not utils.same_shape(value_grad.shape, value.shape):
+            value_grad = sum_to(value_grad, value.shape)
+        put_grad(value, value_grad)
+
+    return fwd
+
+
+register_grad(pids.SETITEM, _setitem_grad)
+
+
 def _log_sigmoid_grad(
     a: TensorProxy,
 ) -> TensorProxy:

From ee050f2f7562778a2d22693a8fc066316fcde285 Mon Sep 17 00:00:00 2001
From: Masato Shinokawa <shino16@users.noreply.github.com>
Date: Tue, 9 Dec 2025 23:33:07 -0800
Subject: [PATCH 04/15] Update testsAdd backward

---
 thunder/core/transform_common.py |  3 ++
 thunder/core/transforms.py       | 51 +++++++++++++++++++-------------
 2 files changed, 34 insertions(+), 20 deletions(-)

diff --git a/thunder/core/transform_common.py b/thunder/core/transform_common.py
index bc14a671f2..9789c136db 100644
--- a/thunder/core/transform_common.py
+++ b/thunder/core/transform_common.py
@@ -235,6 +235,9 @@ def replace_redundant_inputs(
 # into one for ops in this set.
 NON_FUNCTIONAL_OPS: set[prims.PrimIDs | str] = {
     prims.PrimIDs.UNIFORM,
+    prims.PrimIDs.EMPTY,
+    "empty",
+    "torch.empty",
     "torch.uniform",  # this doesn't exist as of the PR
     "torch.uniform_like",  # this doesn't exist as of the PR
     # thunder.core.prims doesn't support. See https://pytorch.org/docs/stable/generated/torch.rand.html.
diff --git a/thunder/core/transforms.py b/thunder/core/transforms.py
index d1a247cd05..b680cb6737 100644
--- a/thunder/core/transforms.py
+++ b/thunder/core/transforms.py
@@ -1442,26 +1442,6 @@ def _maximum_grad(a: TensorProxy, b: TensorProxy, /):
 register_grad(pids.SHAPE, prims.shape)
 
 
-def _setitem_grad(a: TensorProxy, index, value: Number | TensorProxy):
-    fwd = prims.setitem(a, index, value)
-    g = get_grad(fwd)
-
-    a_grad = prims.setitem(prims.clone(g), index, 0)
-    put_grad(a, a_grad)
-
-    if isinstance(value, TensorProxy):
-        value_grad = g[index]
-        # NOTE: `value` could be broadcasted.
-        if not utils.same_shape(value_grad.shape, value.shape):
-            value_grad = sum_to(value_grad, value.shape)
-        put_grad(value, value_grad)
-
-    return fwd
-
-
-register_grad(pids.SETITEM, _setitem_grad)
-
-
 def _log_sigmoid_grad(
     a: TensorProxy,
 ) -> TensorProxy:
@@ -2594,6 +2574,34 @@ def index_put_aug_fwd(
     return VJPDual(primal, residuals)
 
 
+@register_augmented_forward(prims.PrimIDs.SETITEM)
+def setitem_aug_fwd(a, index, value) -> VJPDual:
+    primal = prims.setitem(a, index, value)
+    value_shape = value.shape if isinstance(value, TensorProxy) else None
+    return VJPDual(primal, (index, value_shape))
+
+
+@register_backward(prims.PrimIDs.SETITEM)
+def setitem_backward(index, value_shape, g):
+    # We avoid using Tensor.clone because nvfuserex has unsoundness in mutation on cloned tensors
+    # See https://github.com/Lightning-AI/lightning-thunder/issues/2793
+    def clone(t):
+        cd = get_compile_data()
+        buf = prims.empty(t.shape, device=t.device, dtype=t.dtype)
+        return prims.copy_(t, buf, grad_enabled=cd.is_grad_enabled if cd is not None else False)
+
+    a_grad = prims.setitem(clone(g), index, 0)
+
+    value_grad = None
+    if value_shape is not None:
+        value_grad = g[index]
+        # NOTE: `value` could be broadcasted.
+        if not utils.same_shape(value_grad.shape, value_shape):
+            value_grad = sum_to(value_grad, value_shape)
+
+    return a_grad, None, value_grad
+
+
 if torch.distributed.is_available():
     from torch.distributed import ReduceOp
     from torch._C._distributed_c10d import _resolve_process_group
@@ -3046,8 +3054,11 @@ def vjp(func):
     """
 
     def _vjp(primals, cotangents, **kwargs):
+        from thunder.core.update_aliases import insert_alias_updates
+
         flat_func, flat_args, spec = flatten_func(func, primals, kwargs)
         trace = construct_trace()(flat_func, *flat_args)
+        trace = insert_alias_updates(trace, [])
         result, vjp_result = vjp_call(flat_args, cotangents, trace=trace)
         # If the argument is a CPU scalar tensor, its gradient needs to be summed into a scalar tensor.
         vjp_result = tuple(

From e63da18d0107c8f9030a768aac2c9d46a239860a Mon Sep 17 00:00:00 2001
From: Masato Shinokawa <shino16@users.noreply.github.com>
Date: Tue, 9 Dec 2025 23:37:54 -0800
Subject: [PATCH 05/15] Add test

---
 thunder/tests/opinfos.py   | 22 ++++++++++++++++++++++
 thunder/tests/test_grad.py | 34 ++++++++++++++++++++++++++++++++++
 thunder/torch/__init__.py  |  7 +++++++
 3 files changed, 63 insertions(+)

diff --git a/thunder/tests/opinfos.py b/thunder/tests/opinfos.py
index 7f06d7eab8..76db59f052 100644
--- a/thunder/tests/opinfos.py
+++ b/thunder/tests/opinfos.py
@@ -4473,6 +4473,28 @@ def make_nd_idx(dim_length: int, indices: int, ndim: int):
 shape_ops.append(getitem_opinfo)
 
 
+def setitem_sample_generator(op, device, dtype, requires_grad, **kwargs):
+    for sample in getitem_sample_generator(op, device, dtype, requires_grad, **kwargs):
+        tensor, key = sample.args
+
+        indexed_tensor = tensor[key]
+        value = make_tensor(indexed_tensor.shape, device=device, dtype=dtype, requires_grad=requires_grad)
+        yield SampleInput(tensor, key, value)
+
+        pre_broadcast_shape = tuple(random.choice((s, 1)) for s in indexed_tensor.shape)
+        value = make_tensor(pre_broadcast_shape, device=device, dtype=dtype, requires_grad=requires_grad)
+        yield SampleInput(tensor, key, value)
+
+
+setitem_opinfo = OpInfo(
+    operator.setitem,
+    sample_input_generator=setitem_sample_generator,
+    torch_reference=operator.setitem,
+    numpy_reference=operator.setitem,
+)
+shape_ops.append(setitem_opinfo)
+
+
 def movedim_sample_generator(op, device, dtype, requires_grad, **kwargs):
     make = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
 
diff --git a/thunder/tests/test_grad.py b/thunder/tests/test_grad.py
index 7e17fb04a4..072b79d560 100644
--- a/thunder/tests/test_grad.py
+++ b/thunder/tests/test_grad.py
@@ -42,6 +42,7 @@
     "index_select",
     # Finite difference approximation doesn't work for this function
     "embedding",
+    "setitem",
     "index_put",
     "batch_norm",
     "instance_norm",
@@ -689,6 +690,39 @@ def test_vjp_correctness_embedding_manual(op, device, dtype, executor, comp):
         comp(actual_out, out)
 
 
+@ops((get_opinfo("setitem"),), supported_dtypes=(dtypes.float64,))
+def test_vjp_correctness_setitem_manual(op, device, dtype, executor, comp):
+    for sample in op.sample_inputs(device, dtype, requires_grad=True):
+
+        def torch_reference(tensor, idx, value):
+            cloned = tensor * 1
+            op.torch_reference(cloned, idx, value)
+            return cloned
+
+        def op_fn(tensor, idx, value):
+            cloned = tensor * 1
+            op.op(cloned, idx, value)
+            return cloned
+
+        args_ref = (sample.args[0].detach().clone().requires_grad_(True),) + sample.args[1:]
+        out = torch_reference(*args_ref, **sample.kwargs)
+        v = make_tensor_like(out)
+        expected = torch.autograd.grad(out, (args_ref[0], args_ref[2]), v)
+
+        # Compute vjp result using Thunder
+        flat_op, flat_args, spec = flatten_func(op_fn, sample.args, sample.kwargs)
+        initial_trace = thunder.trace()(vjp(flat_op), flat_args, (v,))
+        jfn = executor.make_callable(initial_trace.python_callable(), disable_torch_autograd=True)
+        actual_out, actual_grad = jfn(flat_args, (v,))
+
+        # With advanced indexing, an element may be assigned multiple times and the assignment order is not guaranteed.
+        # comp(actual_out, out)
+
+        comp(sample.args[0], args_ref[0])
+        comp(actual_grad[0], expected[0])
+        comp(actual_grad[-1], expected[1])
+
+
 @ops((op for op in opinfos if op.name == "type_as"), supported_dtypes=(dtypes.float64,))
 def test_vjp_correctness_type_as_manual(op, device, dtype, executor, comp):
     for sample in op.sample_inputs(device, dtype, requires_grad=True):
diff --git a/thunder/torch/__init__.py b/thunder/torch/__init__.py
index 394f823f8f..d9130c0c18 100644
--- a/thunder/torch/__init__.py
+++ b/thunder/torch/__init__.py
@@ -269,6 +269,13 @@ def _copy_(a, b, /):
     return prims.copy_(b, a, grad_enabled=cd.is_grad_enabled if cd is not None else False)
 
 
+def _clone_via_copy(t: TensorProxy) -> TensorProxy:
+    """Produces a functional clone using an explicit copy instead of prims.clone."""
+    cd = get_compile_data()
+    buf = prims.empty(t.shape, device=t.device, dtype=t.dtype)
+    return prims.copy_(t, buf, grad_enabled=cd.is_grad_enabled if cd is not None else False)
+
+
 @torchsymbol(torch.Tensor.copy_, is_method=True)  # , tags=(prims.OpTags.IN_PLACE,))
 def copy_(a, b, /):
     return _copy_(a, b)

From c98bfcc8aa30c1884f885a7331215f8b83f3685a Mon Sep 17 00:00:00 2001
From: Masato Shinokawa <shino16@users.noreply.github.com>
Date: Wed, 10 Dec 2025 08:47:53 -0800
Subject: [PATCH 06/15] Reduce setitem sample args

---
 thunder/tests/opinfos.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/thunder/tests/opinfos.py b/thunder/tests/opinfos.py
index 76db59f052..54ddba2d28 100644
--- a/thunder/tests/opinfos.py
+++ b/thunder/tests/opinfos.py
@@ -4478,8 +4478,9 @@ def setitem_sample_generator(op, device, dtype, requires_grad, **kwargs):
         tensor, key = sample.args
 
         indexed_tensor = tensor[key]
-        value = make_tensor(indexed_tensor.shape, device=device, dtype=dtype, requires_grad=requires_grad)
-        yield SampleInput(tensor, key, value)
+        # getitem already has lots of cases, and doubling it is too time-consuming
+        # value = make_tensor(indexed_tensor.shape, device=device, dtype=dtype, requires_grad=requires_grad)
+        # yield SampleInput(tensor, key, value)
 
         pre_broadcast_shape = tuple(random.choice((s, 1)) for s in indexed_tensor.shape)
         value = make_tensor(pre_broadcast_shape, device=device, dtype=dtype, requires_grad=requires_grad)

From 3d7fa9c39ad867af005a15ce842cad47d0676194 Mon Sep 17 00:00:00 2001
From: Masato Shinokawa <shino16@users.noreply.github.com>
Date: Wed, 10 Dec 2025 12:01:14 -0800
Subject: [PATCH 07/15] Raise NotImplementedError when backward does so

---
 thunder/clang/__init__.py  | 17 +++++++++++++++++
 thunder/core/transforms.py |  1 +
 2 files changed, 18 insertions(+)

diff --git a/thunder/clang/__init__.py b/thunder/clang/__init__.py
index dd2044f69f..39d2ec9360 100644
--- a/thunder/clang/__init__.py
+++ b/thunder/clang/__init__.py
@@ -13,6 +13,7 @@
 import thunder.clang.utils as clang_utils
 from thunder.core import utils
 from thunder.core.baseutils import run_once
+from thunder.core.compile_data import get_compile_data
 from thunder.core.langctxs import langctx, Languages
 import thunder.core.devices as devices
 import thunder.core.dtypes as dtypes
@@ -649,6 +650,22 @@ def wrap_tensor(t: TensorLike, dim_length: int) -> TensorLike:
 
 @clangop()
 def setitem(a: TensorLike, key, value: TensorLike) -> TensorLike:
+    def is_bool_index(k: Any) -> bool:
+        if isinstance(k, TensorProxy) and dtypes.to_dtype(k) == dtypes.bool8:
+            return True
+        if isinstance(k, Sequence) and any(isinstance(k_i, bool) for k_i in k):
+            return True
+        return False
+
+    key = utils.sequencify(key)
+    if any(is_bool_index(k) for k in key):
+        compile_data = get_compile_data()
+        if compile_data is not None and compile_data.is_grad_enabled and isinstance(value, TensorLike):
+            # See VJP for setitem
+            raise NotImplementedError(
+                "setitem with boolean advanced indexing is not supported when grad is enabled and value is a tensor"
+            )
+
     # TODO: do more checking here. We used to have a check
     #     lambda: f"{key=} tries to index more dimensions than {a.ndim=}",
     return prims.setitem(a, key, value)
diff --git a/thunder/core/transforms.py b/thunder/core/transforms.py
index b680cb6737..8b9ad72fd2 100644
--- a/thunder/core/transforms.py
+++ b/thunder/core/transforms.py
@@ -2594,6 +2594,7 @@ def clone(t):
 
     value_grad = None
     if value_shape is not None:
+        # When index is a boolean tensor, g[index] raises a NotImplementedError because the result shape is unknown
         value_grad = g[index]
         # NOTE: `value` could be broadcasted.
         if not utils.same_shape(value_grad.shape, value_shape):

From 7011cc29a42b92589ccab801eb4c292b65af4ea2 Mon Sep 17 00:00:00 2001
From: Masato Shinokawa <shino16@users.noreply.github.com>
Date: Wed, 10 Dec 2025 14:16:47 -0800
Subject: [PATCH 08/15] Update tests

---
 thunder/tests/opinfos.py   | 37 +++++++++++++++++++++++++++++--------
 thunder/tests/test_grad.py | 23 +++++++++++++++++------
 2 files changed, 46 insertions(+), 14 deletions(-)

diff --git a/thunder/tests/opinfos.py b/thunder/tests/opinfos.py
index 54ddba2d28..c57a15d242 100644
--- a/thunder/tests/opinfos.py
+++ b/thunder/tests/opinfos.py
@@ -4474,17 +4474,38 @@ def make_nd_idx(dim_length: int, indices: int, ndim: int):
 
 
 def setitem_sample_generator(op, device, dtype, requires_grad, **kwargs):
-    for sample in getitem_sample_generator(op, device, dtype, requires_grad, **kwargs):
-        tensor, key = sample.args
+    make = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
 
-        indexed_tensor = tensor[key]
-        # getitem already has lots of cases, and doubling it is too time-consuming
-        # value = make_tensor(indexed_tensor.shape, device=device, dtype=dtype, requires_grad=requires_grad)
+    def _make_setitem_sample(tensor, key):
+        indexed_shape = tensor[key].shape
+
+        # Tests for getitem are already slow, and doubling them is too time-consuming
+        # value = make_tensor(indexed_shape, device=device, dtype=dtype, requires_grad=requires_grad)
         # yield SampleInput(tensor, key, value)
 
-        pre_broadcast_shape = tuple(random.choice((s, 1)) for s in indexed_tensor.shape)
-        value = make_tensor(pre_broadcast_shape, device=device, dtype=dtype, requires_grad=requires_grad)
-        yield SampleInput(tensor, key, value)
+        pre_broadcast_shape = tuple(random.choice((s, 1)) for s in indexed_shape)
+        pre_broadcast_value = make_tensor(pre_broadcast_shape, device=device, dtype=dtype, requires_grad=requires_grad)
+        return SampleInput(tensor, key, pre_broadcast_value)
+
+    for sample in getitem_sample_generator(op, device, dtype, requires_grad, **kwargs):
+        tensor, key = sample.args
+        yield _make_setitem_sample(tensor, key)
+
+    # Boolean mask indexing
+    boolean_mask_cases = [
+        ((6,), (torch.tensor([True, False, True, False, True, False]),)),
+        ((2, 3), (torch.tensor([[True, False, True], [False, True, False]]),)),
+        ((2, 3, 4), ([False, True], [False, True, False], slice(None))),
+        ((2, 3, 4), (torch.tensor([True, False]), [1, 1], slice(None))),
+        ((2, 3, 4), (torch.tensor([False, False]), [1, 1], slice(None))),
+        ((2, 3, 4), (1, torch.tensor([True, False, True]), slice(None))),
+        ((2, 3), (torch.tensor([True, False]), None, [0, 2])),
+        ((4, 2, 3), (Ellipsis, [False, True, False])),
+    ]
+
+    for shape, key in boolean_mask_cases:
+        tensor = make(shape)
+        yield _make_setitem_sample(tensor, key)
 
 
 setitem_opinfo = OpInfo(
diff --git a/thunder/tests/test_grad.py b/thunder/tests/test_grad.py
index 072b79d560..f6eb7b91f1 100644
--- a/thunder/tests/test_grad.py
+++ b/thunder/tests/test_grad.py
@@ -704,13 +704,24 @@ def op_fn(tensor, idx, value):
             op.op(cloned, idx, value)
             return cloned
 
-        args_ref = (sample.args[0].detach().clone().requires_grad_(True),) + sample.args[1:]
-        out = torch_reference(*args_ref, **sample.kwargs)
+        tensor, key, value = sample.args
+        assert not sample.kwargs
+
+        tensor_ref = tensor.detach().clone().requires_grad_(True)
+        out = torch_reference(tensor_ref, key, value)
         v = make_tensor_like(out)
-        expected = torch.autograd.grad(out, (args_ref[0], args_ref[2]), v)
+        expected = torch.autograd.grad(out, (tensor_ref, value), v)
+
+        flat_op, flat_args, spec = flatten_func(op_fn, (tensor, key, value), {})
+
+        t_key = key if isinstance(key, tuple) else (key,)
+        if any(isinstance(k, (torch.Tensor, Sequence)) and torch.tensor(k).dtype == torch.bool for k in t_key):
+            with pytest.raises(NotImplementedError):
+                executor.make_callable(flat_op, disable_torch_autograd=True)(*flat_args)
+            with pytest.raises(NotImplementedError):
+                vjp(flat_op)(flat_args, (v,))
+            continue
 
-        # Compute vjp result using Thunder
-        flat_op, flat_args, spec = flatten_func(op_fn, sample.args, sample.kwargs)
         initial_trace = thunder.trace()(vjp(flat_op), flat_args, (v,))
         jfn = executor.make_callable(initial_trace.python_callable(), disable_torch_autograd=True)
         actual_out, actual_grad = jfn(flat_args, (v,))
@@ -718,7 +729,7 @@ def op_fn(tensor, idx, value):
         # With advanced indexing, an element may be assigned multiple times and the assignment order is not guaranteed.
         # comp(actual_out, out)
 
-        comp(sample.args[0], args_ref[0])
+        comp(tensor, tensor_ref)
         comp(actual_grad[0], expected[0])
         comp(actual_grad[-1], expected[1])
 

From d9a551b1197cd15853e76258586854dffb171b88 Mon Sep 17 00:00:00 2001
From: Masato Shinokawa <shino16@users.noreply.github.com>
Date: Thu, 11 Dec 2025 03:37:33 -0800
Subject: [PATCH 09/15] Revert "Raise NotImplementedError when backward does
 so"

This reverts commit 3d7fa9c39ad867af005a15ce842cad47d0676194.
---
 thunder/clang/__init__.py  | 17 -----------------
 thunder/core/transforms.py |  1 -
 2 files changed, 18 deletions(-)

diff --git a/thunder/clang/__init__.py b/thunder/clang/__init__.py
index 39d2ec9360..dd2044f69f 100644
--- a/thunder/clang/__init__.py
+++ b/thunder/clang/__init__.py
@@ -13,7 +13,6 @@
 import thunder.clang.utils as clang_utils
 from thunder.core import utils
 from thunder.core.baseutils import run_once
-from thunder.core.compile_data import get_compile_data
 from thunder.core.langctxs import langctx, Languages
 import thunder.core.devices as devices
 import thunder.core.dtypes as dtypes
@@ -650,22 +649,6 @@ def wrap_tensor(t: TensorLike, dim_length: int) -> TensorLike:
 
 @clangop()
 def setitem(a: TensorLike, key, value: TensorLike) -> TensorLike:
-    def is_bool_index(k: Any) -> bool:
-        if isinstance(k, TensorProxy) and dtypes.to_dtype(k) == dtypes.bool8:
-            return True
-        if isinstance(k, Sequence) and any(isinstance(k_i, bool) for k_i in k):
-            return True
-        return False
-
-    key = utils.sequencify(key)
-    if any(is_bool_index(k) for k in key):
-        compile_data = get_compile_data()
-        if compile_data is not None and compile_data.is_grad_enabled and isinstance(value, TensorLike):
-            # See VJP for setitem
-            raise NotImplementedError(
-                "setitem with boolean advanced indexing is not supported when grad is enabled and value is a tensor"
-            )
-
     # TODO: do more checking here. We used to have a check
     #     lambda: f"{key=} tries to index more dimensions than {a.ndim=}",
     return prims.setitem(a, key, value)
diff --git a/thunder/core/transforms.py b/thunder/core/transforms.py
index 8b9ad72fd2..b680cb6737 100644
--- a/thunder/core/transforms.py
+++ b/thunder/core/transforms.py
@@ -2594,7 +2594,6 @@ def clone(t):
 
     value_grad = None
     if value_shape is not None:
-        # When index is a boolean tensor, g[index] raises a NotImplementedError because the result shape is unknown
         value_grad = g[index]
         # NOTE: `value` could be broadcasted.
         if not utils.same_shape(value_grad.shape, value_shape):

From 2a42650b8272a1d25c3f814644dc7b82b77515cd Mon Sep 17 00:00:00 2001
From: Masato Shinokawa <shino16@users.noreply.github.com>
Date: Thu, 11 Dec 2025 07:42:12 -0800
Subject: [PATCH 10/15] Fix opinfos sample input

---
 thunder/tests/opinfos.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/thunder/tests/opinfos.py b/thunder/tests/opinfos.py
index c57a15d242..dc79d98072 100644
--- a/thunder/tests/opinfos.py
+++ b/thunder/tests/opinfos.py
@@ -4496,8 +4496,8 @@ def _make_setitem_sample(tensor, key):
         ((6,), (torch.tensor([True, False, True, False, True, False]),)),
         ((2, 3), (torch.tensor([[True, False, True], [False, True, False]]),)),
         ((2, 3, 4), ([False, True], [False, True, False], slice(None))),
+        ((2, 3, 4), (torch.tensor([False, False]), slice(None))),
         ((2, 3, 4), (torch.tensor([True, False]), [1, 1], slice(None))),
-        ((2, 3, 4), (torch.tensor([False, False]), [1, 1], slice(None))),
         ((2, 3, 4), (1, torch.tensor([True, False, True]), slice(None))),
         ((2, 3), (torch.tensor([True, False]), None, [0, 2])),
         ((4, 2, 3), (Ellipsis, [False, True, False])),

From 8a08ba4f294293ceca09b292529c966455f93bdf Mon Sep 17 00:00:00 2001
From: Masato Shinokawa <shino16@users.noreply.github.com>
Date: Thu, 11 Dec 2025 07:46:00 -0800
Subject: [PATCH 11/15] Make backward of value_and_grad(setitem) non-empty

---
 thunder/dynamo/utils.py | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/thunder/dynamo/utils.py b/thunder/dynamo/utils.py
index 8c643cee71..783c4b7fa0 100644
--- a/thunder/dynamo/utils.py
+++ b/thunder/dynamo/utils.py
@@ -2,6 +2,7 @@
 from collections.abc import Callable, Sequence
 from contextlib import contextmanager
 from enum import Enum, auto
+import operator
 from typing import TYPE_CHECKING
 import dataclasses
 import inspect
@@ -389,11 +390,19 @@ def _run_with_cache_info():
                 exception=str(e),
             )
 
-        function_to_run = (
-            value_and_grad(thunder_symbol)
-            if requires_grad and (disable_torch_autograd is None or not disable_torch_autograd)
-            else thunder_symbol
-        )
+        if requires_grad and (disable_torch_autograd is None or not disable_torch_autograd):
+            if thunder_symbol is operator.setitem:
+                # operator.setitem returns None, which makes its backward pass empty
+                # We don't need to cover torch.Tensor.__setitem__ as dynamo uses operator.setitem instead
+                def setitem_and_return(a, key, value):
+                    a[key] = value
+                    return a
+
+                function_to_run = value_and_grad(setitem_and_return)
+            else:
+                function_to_run = value_and_grad(thunder_symbol)
+        else:
+            function_to_run = thunder_symbol
         # We need to be under trace context to generate proxies.
         with thunder.core.trace.tracectx(TraceCtx()):
             try:

From d7ca6a8ad9f1de31bd7bc8bb085ca9006e396610 Mon Sep 17 00:00:00 2001
From: Masato Shinokawa <shino16@users.noreply.github.com>
Date: Thu, 11 Dec 2025 07:46:04 -0800
Subject: [PATCH 12/15] Update test

---
 thunder/tests/test_grad.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/thunder/tests/test_grad.py b/thunder/tests/test_grad.py
index f6eb7b91f1..66b818df8b 100644
--- a/thunder/tests/test_grad.py
+++ b/thunder/tests/test_grad.py
@@ -716,8 +716,6 @@ def op_fn(tensor, idx, value):
 
         t_key = key if isinstance(key, tuple) else (key,)
         if any(isinstance(k, (torch.Tensor, Sequence)) and torch.tensor(k).dtype == torch.bool for k in t_key):
-            with pytest.raises(NotImplementedError):
-                executor.make_callable(flat_op, disable_torch_autograd=True)(*flat_args)
             with pytest.raises(NotImplementedError):
                 vjp(flat_op)(flat_args, (v,))
             continue

From 884109c8a69be1ce189b1c2f2d2169ce95635fcd Mon Sep 17 00:00:00 2001
From: Masato Shinokawa <shino16@users.noreply.github.com>
Date: Thu, 11 Dec 2025 09:35:12 -0800
Subject: [PATCH 13/15] Add test

---
 thunder/core/transforms.py |  5 +++--
 thunder/tests/test_grad.py | 39 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/thunder/core/transforms.py b/thunder/core/transforms.py
index b680cb6737..c996808a3d 100644
--- a/thunder/core/transforms.py
+++ b/thunder/core/transforms.py
@@ -12,6 +12,7 @@
 import time
 import dataclasses
 
+from thunder.core.update_aliases import insert_alias_updates
 import thunder.core.utils as utils
 from thunder.core import dtypes, prims
 from thunder.core.devices import Device
@@ -3054,10 +3055,10 @@ def vjp(func):
     """
 
     def _vjp(primals, cotangents, **kwargs):
-        from thunder.core.update_aliases import insert_alias_updates
-
         flat_func, flat_args, spec = flatten_func(func, primals, kwargs)
         trace = construct_trace()(flat_func, *flat_args)
+        # Although we do not need to insert prims.update_aliases, we need insert_alias_updates for variable substitution
+        # e.g. `a.sin_(); return a` into `b = a.sin_(); return b`
         trace = insert_alias_updates(trace, [])
         result, vjp_result = vjp_call(flat_args, cotangents, trace=trace)
         # If the argument is a CPU scalar tensor, its gradient needs to be summed into a scalar tensor.
diff --git a/thunder/tests/test_grad.py b/thunder/tests/test_grad.py
index 66b818df8b..826090acf3 100644
--- a/thunder/tests/test_grad.py
+++ b/thunder/tests/test_grad.py
@@ -11,12 +11,14 @@
 import torch
 
 import thunder
+from thunder.core import prims
 import thunder.core.dtypes as dtypes
 import thunder.core.devices as devices
 
 from thunder import torch as ltorch
 from thunder.core.dtypes import is_exact_dtype, to_dtype as thunder_dtype
 from thunder.core.pytree import tree_map, tree_flatten
+from thunder.core.symbol import BoundSymbol
 from thunder.core.transforms import vjp, grad, check_bsym_for_vjp
 from thunder.core.utils import flatten_func, is_cpu_scalar_tensor
 from thunder.tests.framework import (
@@ -995,6 +997,43 @@ def test_vjp_correctness_einsum_manual(op, device, dtype, executor, comp):
             comp(torch_grad, thunder_grad)
 
 
+@ops((get_opinfo("sin"), get_opinfo("mul")), supported_dtypes=(dtypes.float64,))
+def test_vjp_correctness_inplace(op, device, dtype, executor, comp):
+    if op.op.__name__ == "sin":
+        fn = lambda a: a.clone().sin_()
+    else:
+        fn = lambda a, b: a.clone().mul_(b)
+
+    for sample in op.sample_inputs(device, dtype, requires_grad=True):
+        assert not sample.kwargs
+
+        ref_args = (sample.args[0].clone().detach().requires_grad_(sample.args[0].requires_grad),) + sample.args[1:]
+        out = fn(*ref_args)
+        v = make_tensor_like(out)
+        expected_grads = torch.autograd.grad(out, ref_args, v)
+
+        initial_trace = thunder.trace()(vjp(fn), sample.args, (v,))
+
+        def map_prim_copy_to_ltorch_copy(bsym: BoundSymbol):
+            if bsym.sym.id == prims.PrimIDs.COPY_:
+                return bsym.from_bsym(
+                    sym=thunder.torch.copy_, args=(bsym.args[1], bsym.args[0]), kwargs={}, subsymbols=[bsym]
+                )
+            else:
+                return bsym
+
+        initial_trace.bound_symbols = list(map(map_prim_copy_to_ltorch_copy, initial_trace.bound_symbols))
+
+        actual_out, grads_out = executor.make_callable(initial_trace.python_callable(), disable_torch_autograd=True)(
+            sample.args, (v,)
+        )
+
+        comp(actual_out, out)
+        comp(ref_args[0], sample.args[0])
+        for torch_grad, thunder_grad in zip(expected_grads, grads_out):
+            comp(torch_grad, thunder_grad)
+
+
 # TODO Extend requires_grad so that tensors produced from thunder.jit functions requires_grad
 #   and have their autograd functions set properly
 # Tests that we track the requires_grad property properly

From a8f276b9a369aac77fb424727488de54c5ad0ed9 Mon Sep 17 00:00:00 2001
From: Masato Shinokawa <shino16@users.noreply.github.com>
Date: Thu, 11 Dec 2025 09:42:01 -0800
Subject: [PATCH 14/15] Revert "Add test"

This reverts commit 884109c8a69be1ce189b1c2f2d2169ce95635fcd.
---
 thunder/core/transforms.py |  5 ++---
 thunder/tests/test_grad.py | 39 --------------------------------------
 2 files changed, 2 insertions(+), 42 deletions(-)

diff --git a/thunder/core/transforms.py b/thunder/core/transforms.py
index c996808a3d..b680cb6737 100644
--- a/thunder/core/transforms.py
+++ b/thunder/core/transforms.py
@@ -12,7 +12,6 @@
 import time
 import dataclasses
 
-from thunder.core.update_aliases import insert_alias_updates
 import thunder.core.utils as utils
 from thunder.core import dtypes, prims
 from thunder.core.devices import Device
@@ -3055,10 +3054,10 @@ def vjp(func):
     """
 
     def _vjp(primals, cotangents, **kwargs):
+        from thunder.core.update_aliases import insert_alias_updates
+
         flat_func, flat_args, spec = flatten_func(func, primals, kwargs)
         trace = construct_trace()(flat_func, *flat_args)
-        # Although we do not need to insert prims.update_aliases, we need insert_alias_updates for variable substitution
-        # e.g. `a.sin_(); return a` into `b = a.sin_(); return b`
         trace = insert_alias_updates(trace, [])
         result, vjp_result = vjp_call(flat_args, cotangents, trace=trace)
         # If the argument is a CPU scalar tensor, its gradient needs to be summed into a scalar tensor.
diff --git a/thunder/tests/test_grad.py b/thunder/tests/test_grad.py
index 826090acf3..66b818df8b 100644
--- a/thunder/tests/test_grad.py
+++ b/thunder/tests/test_grad.py
@@ -11,14 +11,12 @@
 import torch
 
 import thunder
-from thunder.core import prims
 import thunder.core.dtypes as dtypes
 import thunder.core.devices as devices
 
 from thunder import torch as ltorch
 from thunder.core.dtypes import is_exact_dtype, to_dtype as thunder_dtype
 from thunder.core.pytree import tree_map, tree_flatten
-from thunder.core.symbol import BoundSymbol
 from thunder.core.transforms import vjp, grad, check_bsym_for_vjp
 from thunder.core.utils import flatten_func, is_cpu_scalar_tensor
 from thunder.tests.framework import (
@@ -997,43 +995,6 @@ def test_vjp_correctness_einsum_manual(op, device, dtype, executor, comp):
             comp(torch_grad, thunder_grad)
 
 
-@ops((get_opinfo("sin"), get_opinfo("mul")), supported_dtypes=(dtypes.float64,))
-def test_vjp_correctness_inplace(op, device, dtype, executor, comp):
-    if op.op.__name__ == "sin":
-        fn = lambda a: a.clone().sin_()
-    else:
-        fn = lambda a, b: a.clone().mul_(b)
-
-    for sample in op.sample_inputs(device, dtype, requires_grad=True):
-        assert not sample.kwargs
-
-        ref_args = (sample.args[0].clone().detach().requires_grad_(sample.args[0].requires_grad),) + sample.args[1:]
-        out = fn(*ref_args)
-        v = make_tensor_like(out)
-        expected_grads = torch.autograd.grad(out, ref_args, v)
-
-        initial_trace = thunder.trace()(vjp(fn), sample.args, (v,))
-
-        def map_prim_copy_to_ltorch_copy(bsym: BoundSymbol):
-            if bsym.sym.id == prims.PrimIDs.COPY_:
-                return bsym.from_bsym(
-                    sym=thunder.torch.copy_, args=(bsym.args[1], bsym.args[0]), kwargs={}, subsymbols=[bsym]
-                )
-            else:
-                return bsym
-
-        initial_trace.bound_symbols = list(map(map_prim_copy_to_ltorch_copy, initial_trace.bound_symbols))
-
-        actual_out, grads_out = executor.make_callable(initial_trace.python_callable(), disable_torch_autograd=True)(
-            sample.args, (v,)
-        )
-
-        comp(actual_out, out)
-        comp(ref_args[0], sample.args[0])
-        for torch_grad, thunder_grad in zip(expected_grads, grads_out):
-            comp(torch_grad, thunder_grad)
-
-
 # TODO Extend requires_grad so that tensors produced from thunder.jit functions requires_grad
 #   and have their autograd functions set properly
 # Tests that we track the requires_grad property properly

From a64f0ba4f2d67cc6ea9321fce902db2dbc07df1d Mon Sep 17 00:00:00 2001
From: Masato Shinokawa <shino16@users.noreply.github.com>
Date: Thu, 11 Dec 2025 09:44:11 -0800
Subject: [PATCH 15/15] Cosmetic change

---
 thunder/core/transforms.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/thunder/core/transforms.py b/thunder/core/transforms.py
index b680cb6737..aa5b45df3c 100644
--- a/thunder/core/transforms.py
+++ b/thunder/core/transforms.py
@@ -12,6 +12,7 @@
 import time
 import dataclasses
 
+from thunder.core.update_aliases import insert_alias_updates
 import thunder.core.utils as utils
 from thunder.core import dtypes, prims
 from thunder.core.devices import Device
@@ -3054,10 +3055,9 @@ def vjp(func):
     """
 
     def _vjp(primals, cotangents, **kwargs):
-        from thunder.core.update_aliases import insert_alias_updates
-
         flat_func, flat_args, spec = flatten_func(func, primals, kwargs)
         trace = construct_trace()(flat_func, *flat_args)
+        # No need to insert prims.update_aliases, but we need insert_alias_updates for variable substitution
         trace = insert_alias_updates(trace, [])
         result, vjp_result = vjp_call(flat_args, cotangents, trace=trace)
         # If the argument is a CPU scalar tensor, its gradient needs to be summed into a scalar tensor.