Enable interpolate VJP tests, add PrimID mapping for ceil and floor

Author: aobolensk

thunder/core/transforms.py

@@ -1660,6 +1660,8 @@ def zeros_like(x):
     prims.PrimIDs.COPY_: lambda x, y, grad_enabled: (prims.copy_(x, y, grad_enabled=grad_enabled), tuple()),
     prims.PrimIDs.CLONE: lambda x: (prims.clone(x), tuple()),
     prims.PrimIDs.BITCAST: lambda x, dtype: (prims.bitcast(x, dtype), (x.dtype,)),
+    prims.PrimIDs.CEIL: lambda x: (prims.ceil(x), (x,)),
+    prims.PrimIDs.FLOOR: lambda x: (prims.floor(x), (x,)),
 }
 
 
@@ -1691,6 +1693,8 @@ def zeros_like(x):
     prims.PrimIDs.COPY_: lambda g: (g, None),
     prims.PrimIDs.CLONE: lambda g: g,
     prims.PrimIDs.BITCAST: lambda x_dtype, g: (prims.bitcast(g, x_dtype), None),
+    prims.PrimIDs.CEIL: lambda x, g: g * 0,
+    prims.PrimIDs.FLOOR: lambda x, g: g * 0,
 }
 
 

thunder/tests/opinfos.py

@@ -9839,20 +9839,6 @@ def interpolate_error_generator(op, device, dtype=torch.float32, **kwargs):
     error_input_generator=interpolate_error_generator,
     torch_reference=torch.nn.functional.interpolate,
     dtypes=(datatypes.floating,),
-    test_directives=(
-        # PyTorch does not support CPU Half upsample used in interpolate
-        DecorateInfo(
-            pytest.mark.xfail,
-            "test_core_vs_torch_consistency",
-            dtypes=(datatypes.float16,),
-            devicetypes=(devices.DeviceType.CPU,),
-        ),
-        # This should be fixed now; TODO re-enable and test
-        DecorateInfo(
-            pytest.mark.xfail,
-            "test_vjp_correctness",
-        ),
-    ),
 )
 nn_ops.append(interpolate_opinfo)