remove flattening and unflattening

of tensor subclasses from `nondifferentiable_vjp_symbols` since the trace transform of tensor subclasses comes after VJP Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
Lightning-AI · Dec 23, 2024 · 56c69df · 56c69df
1 parent e5d26fc
commit 56c69df
Showing 1 changed file with 0 additions and 2 deletions.
diff --git a/thunder/core/transforms.py b/thunder/core/transforms.py
@@ -2509,8 +2509,6 @@ def uniform_backward(primal, minval, maxval, g):
     prims.PrimIDs.BITWISE_XOR,
     prims.PrimIDs.SIGNBIT,
     prims.PrimIDs.FULL,
-    prims.PrimIDs.FLATTEN_TENSOR_SUBCLASS,
-    prims.PrimIDs.UNFLATTEN_TENSOR_SUBCLASS,
 }