improve check

Signed-off-by: Masaki Kozuki <[email protected]>

Author: crcrpar

thunder/executors/cutlass_dsl_ex.py

@@ -105,10 +105,10 @@ def is_last_dim_divisible(dtype: dtypes.dtype, last_dim_size: int) -> bool:
 
     def quack_softmax_impl(a: torch.Tensor) -> torch.Tensor:
         original_shape = a.shape
-        if requires_reshpae := a.ndim > 2:
+        if requires_reshape := a.ndim > 2:
             a = a.view(-1, original_shape[-1])
         ret = softmax_fwd(a)
-        if requires_reshpae:
+        if requires_reshape:
             ret = ret.view(original_shape)
         return ret
 
@@ -273,11 +273,11 @@ def quack_cross_entropy_checker(
         if label_smoothing != 0.0:
             return False
 
-        if (
-            a.ndim != 2
-            or a.dtype not in {dtypes.float16, dtypes.bfloat16, dtypes.float32}
+        if not (
+            a.ndim == 2
+            and a.dtype in {dtypes.float16, dtypes.bfloat16, dtypes.float32}
             and target.ndim == 1
-            and target.dytpe in {dtypes.int32, dtypes.int64}
+            and target.dtype in {dtypes.int32, dtypes.int64}
         ):
             return False
 
@@ -383,7 +383,7 @@ def quack_layer_norm_checker(
     ) -> bool:
         if (
             a.dtype not in {dtypes.float16, dtypes.bfloat16, dtypes.float32}
-            or weight.ndim != 1
+            or (weight is None or weight.ndim != 1)
             or a.shape[-1] != weight.shape[0]
             or weight.dtype not in {dtypes.float32}
         ):
@@ -473,7 +473,7 @@ def quack_rms_norm_checker(
         eps: float | None = None,
     ) -> bool:
         if (
-            weight.ndim != 1
+            (weight is None or weight.ndim != 1)
             or a.shape[-1] != weight.shape[0]
             or a.dtype not in {dtypes.float16, dtypes.bfloat16, dtypes.float32}
             or weight.dtype not in {dtypes.float16, dtypes.bfloat16, dtypes.float32}