Keep up with NineToothed 0.16.0 updates

Author: voltjia

attention.py

@@ -56,7 +56,10 @@ def application(q, k, v, o):
     o = acc  # noqa: F841
 
 
-q, k, v, o = (Tensor(4, constexpr_shape=True) for _ in range(4))
+q, k, v, o = (
+    Tensor(4, shape_options=(None, None, None, {"constexpr": True, "upper_bound": 128}))
+    for _ in range(4)
+)
 attention_kernel = ninetoothed.make(arrangement, application, (q, k, v, o))
 
 

conv2d.py

@@ -23,11 +23,14 @@ def arrangement(input, filter, output):
     return matmul.arrangement(input_flattened, filter_permuted, output_flattened)
 
 
-conv2d_kernel = ninetoothed.make(
-    arrangement,
-    matmul.application,
-    (Tensor(4), Tensor(4, constexpr_shape=True), Tensor(4)),
+filter_shape_options = (
+    None,
+    None,
+    {"constexpr": True, "upper_bound": 16},
+    {"constexpr": True, "upper_bound": 16},
 )
+tensors = (Tensor(4), Tensor(4, shape_options=filter_shape_options), Tensor(4))
+conv2d_kernel = ninetoothed.make(arrangement, matmul.application, tensors)
 
 
 def conv2d(input, filter):

max_pool2d.py

@@ -11,8 +11,8 @@
 def arrangement(input, output):
     BLOCK_SIZE = Symbol("BLOCK_SIZE", meta=True)
 
-    WINDOW_HEIGHT = Symbol("WINDOW_HEIGHT", constexpr=True)
-    WINDOW_WIDTH = Symbol("WINDOW_WIDTH", constexpr=True)
+    WINDOW_HEIGHT = Symbol("WINDOW_HEIGHT", constexpr=True, upper_bound=16)
+    WINDOW_WIDTH = Symbol("WINDOW_WIDTH", constexpr=True, upper_bound=16)
 
     input_arranged = input.tile((1, 1, WINDOW_HEIGHT, WINDOW_WIDTH))
     input_arranged = input_arranged.ravel()