Fix the dtype values of zero-dimensional tensors in the premake functions

Author: voltjia

src/ntops/kernels/add.py

@@ -1,5 +1,6 @@
 import functools
 
+import ninetoothed
 from ninetoothed import Tensor
 
 from ntops.kernels.element_wise import arrangement
@@ -15,7 +16,7 @@ def premake(ndim, dtype=None, block_size=None):
     tensors = (
         Tensor(ndim, dtype=dtype),
         Tensor(ndim, dtype=dtype),
-        Tensor(0, dtype=dtype),
+        Tensor(0, dtype=ninetoothed.float64),
         Tensor(ndim, dtype=dtype),
     )
 

src/ntops/kernels/addmm.py

@@ -1,5 +1,6 @@
 import functools
 
+import ninetoothed
 import ninetoothed.language as ntl
 from ninetoothed import Tensor
 
@@ -84,10 +85,10 @@ def premake(
         Tensor(2, dtype=dtype),
         Tensor(2, dtype=dtype),
         Tensor(2, dtype=dtype),
-        Tensor(0, dtype=dtype),
-        Tensor(0, dtype=dtype),
+        Tensor(0, dtype=ninetoothed.float64),
+        Tensor(0, dtype=ninetoothed.float64),
         Tensor(2, dtype=dtype),
-        Tensor(0, dtype=dtype, constexpr=True, value=input_precision),
+        Tensor(0, constexpr=True, value=input_precision),
     )
 
     return arrangement_, application, tensors

src/ntops/kernels/bmm.py

@@ -61,7 +61,7 @@ def premake(
         Tensor(3, dtype=dtype),
         Tensor(3, dtype=dtype),
         Tensor(3, dtype=dtype),
-        Tensor(0, dtype=dtype, constexpr=True, value=input_precision),
+        Tensor(0, constexpr=True, value=input_precision),
     )
 
     return arrangement_, application, tensors

src/ntops/kernels/dropout.py

@@ -1,5 +1,6 @@
 import functools
 
+import ninetoothed
 import ninetoothed.language as ntl
 from ninetoothed import Tensor
 
@@ -15,8 +16,8 @@ def premake(ndim, dtype=None, block_size=None):
 
     tensors = (
         Tensor(ndim, dtype=dtype),
-        Tensor(0, dtype=dtype),
-        Tensor(0, dtype=dtype),
+        Tensor(0, dtype=ninetoothed.float64),
+        Tensor(0, dtype=ninetoothed.int64),
         Tensor(ndim, dtype=dtype),
     )
 

src/ntops/kernels/layer_norm.py

@@ -1,6 +1,7 @@
 import functools
 import math
 
+import ninetoothed
 import ninetoothed.language as ntl
 from ninetoothed import Tensor
 
@@ -39,9 +40,9 @@ def premake(ndim, normalized_shape, dtype=None, block_size=None):
         Tensor(ndim, other=0, dtype=dtype),
         Tensor(ndim, dtype=dtype),
         Tensor(ndim, dtype=dtype),
-        Tensor(0, dtype=dtype),
+        Tensor(0, dtype=ninetoothed.float64),
         Tensor(ndim, dtype=dtype),
-        Tensor(0, dtype=dtype, constexpr=True, value=math.prod(normalized_shape)),
+        Tensor(0, constexpr=True, value=math.prod(normalized_shape)),
     )
 
     return arrangement_, application, tensors

src/ntops/kernels/mm.py

@@ -83,7 +83,7 @@ def premake(
         Tensor(2, dtype=dtype),
         Tensor(2, dtype=dtype),
         Tensor(2, dtype=dtype),
-        Tensor(0, dtype=dtype, constexpr=True, value=input_precision),
+        Tensor(0, constexpr=True, value=input_precision),
     )
 
     return arrangement_, application, tensors

src/ntops/kernels/rms_norm.py

@@ -35,9 +35,9 @@ def premake(
     tensors = (
         Tensor(ndim, other=0, dtype=input_dtype),
         Tensor(ndim, dtype=weight_dtype),
-        Tensor(0, dtype=ninetoothed.float32),
+        Tensor(0, dtype=ninetoothed.float64),
         Tensor(ndim, dtype=output_dtype),
-        Tensor(0, dtype=ninetoothed.uint64),
+        Tensor(0, dtype=ninetoothed.int64),
     )
 
     return arrangement_, application, tensors

src/ntops/kernels/scaled_dot_product_attention.py

@@ -229,10 +229,10 @@ def premake(
     present_key, present_value, present_key_slot, present_value_slot = (
         Tensor(4, dtype=dtype) for _ in range(4)
     )
-    scale = Tensor(0, dtype=dtype)
-    is_causal = Tensor(0, dtype=dtype, constexpr=True, value=is_causal)
-    with_attn_mask = Tensor(0, dtype=dtype, constexpr=True, value=with_attn_mask)
-    causal_variant = Tensor(0, dtype=dtype, constexpr=True, value=causal_variant)
+    scale = Tensor(0, dtype=ninetoothed.float64)
+    is_causal = Tensor(0, constexpr=True, value=is_causal)
+    with_attn_mask = Tensor(0, constexpr=True, value=with_attn_mask)
+    causal_variant = Tensor(0, constexpr=True, value=causal_variant)
 
     if emb_dim is not None:
         for tensor in (query, key, value, attn_mask, output):

src/ntops/kernels/sub.py

@@ -1,5 +1,6 @@
 import functools
 
+import ninetoothed
 from ninetoothed import Tensor
 
 from ntops.kernels.element_wise import arrangement
@@ -15,7 +16,7 @@ def premake(ndim, dtype=None, block_size=None):
     tensors = (
         Tensor(ndim, dtype=dtype),
         Tensor(ndim, dtype=dtype),
-        Tensor(0, dtype=dtype),
+        Tensor(0, dtype=ninetoothed.float64),
         Tensor(ndim, dtype=dtype),
     )