Add pad tests (Lightning-AI#536)

rdspring1 · Jun 6, 2024 · d3c06ff · d3c06ff
1 parent 0c6e38c
commit d3c06ff
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/thunder/tests/opinfos.py b/thunder/tests/opinfos.py
@@ -3748,6 +3748,11 @@ def pad_torch_sample_generator(op, device, dtype, requires_grad, **kwargs):
     for shape, padding_config in cases:
         yield SampleInput(make(shape), padding_config, "constant", make_number(dtype=dtype))
 
+    # The `value` parameter of the pad op is unceremoniously cast to the type of the
+    # tensor being padded. Yield some tests with explicitly-different data types.
+    yield SampleInput(make((2, 3)), pad=(1, 2), value=6.4)
+    yield SampleInput(make((2,)), pad=(1, 2), value=1)
+
 
 def pad_torch_error_generator(op, device, dtype=torch.float32, **kwargs):
     make = partial(make_tensor, device=device, dtype=dtype)