Update tests

Author: shino16

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(

thunder/tests/test_grad.py

@@ -704,21 +704,32 @@ 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,))
 
         # 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])