[XLA:CPU] Use distinct element values in tiled kernel tests

Using random values is not best practice.
Using inputs made out of the same elements isn't a good test for ops like transpose.

PiperOrigin-RevId: 837188191

Author: basioli-k

xla/backends/cpu/codegen/tiled/tiled_kernel_test.py

@@ -28,9 +28,13 @@
 
 class InputSpec:
 
-  def __init__(self, shape: tuple[int, ...], input_value):
+  def __init__(self, shape: tuple[int, ...]):
+    """Initializes the InputSpec.
+
+    Args:
+      shape: The shape of the input array.
+    """
     self.shape = shape
-    self.input_value = input_value
 
 
 def get_random_array(shape: tuple[int, ...], dtype: np.dtype) -> np.ndarray:
@@ -58,12 +62,8 @@ def compare_kernel(
       cpu_testlib.JitCompiler(base_testlib.HloModuleConfig()),
   )
 
-  # Simply use a all-ones arrays as inputs to make it easy to debug the kernel
-  # unless random inputs are requested.
   def get_input(spec: InputSpec):
-    if spec.input_value is None:
-      return get_random_array(spec.shape, dtype)
-    return np.full(shape=spec.shape, fill_value=spec.input_value, dtype=dtype)
+    return np.arange(np.prod(spec.shape), dtype=dtype).reshape(spec.shape)
 
   inputs = [get_input(spec) for spec in input_specs]
 
@@ -105,7 +105,7 @@ def test_slice(self):
         ir,
         "tiled_slice",
         1,
-        [InputSpec((5, 5), 1)],
+        [InputSpec((5, 5))],
         (5, 5),
         np.float32,
         lambda arg: arg.transpose(),
@@ -129,7 +129,7 @@ def test_strided(self):
         ir,
         "tiled_slice",
         1,
-        [InputSpec((64, 64), 1)],
+        [InputSpec((64, 64))],
         (4, 32),
         np.float32,
         lambda arg: arg[::21, ::2],
@@ -156,7 +156,7 @@ def test_transpose(self):
         ir,
         "tiled_transpose",
         8,
-        [InputSpec((4096, 4096), 1)],
+        [InputSpec((4096, 4096))],
         (4096, 4096),
         np.float32,
         lambda arg: arg.transpose(),
@@ -185,7 +185,7 @@ def test_add_tranpose(self):
         ir,
         "add_tranpose",
         8,
-        [InputSpec((4096, 4096), 1)],
+        [InputSpec((4096, 4096))],
         (4096, 4096),
         np.float32,
         lambda arg: arg + arg.transpose(),
@@ -213,7 +213,7 @@ def test_dot_single_tile(self):
         ir,
         "dot_single_tile",
         1,
-        [InputSpec((8, 16), 1), InputSpec((16, 8), 1)],
+        [InputSpec((8, 16)), InputSpec((16, 8))],
         (8, 8),
         np.float32,
         lambda lhs, rhs: lhs @ rhs,
@@ -242,7 +242,7 @@ def test_dot_scalar_output(self):
         ir,
         "test_dot_scalar_output",
         1,
-        [InputSpec((8, 16), 1), InputSpec((16, 8), 1)],
+        [InputSpec((8, 16)), InputSpec((16, 8))],
         (),
         np.float32,
         lambda lhs, rhs: np.tensordot(lhs, rhs, axes=[[1, 0], [0, 1]]),
@@ -275,7 +275,11 @@ def test_dot_fusion_single_tile(self):
         ir,
         "dot_fusion_single_tile",
         1,
-        [InputSpec((8, 16), 1), InputSpec((8, 16), 1), InputSpec((16, 1), 1)],
+        [
+            InputSpec((8, 16)),
+            InputSpec((8, 16)),
+            InputSpec((16, 1)),
+        ],
         (8, 1),
         np.float32,
         lambda lhs_0, lhs_1, rhs: np.tanh((lhs_0 + lhs_1) @ rhs),
@@ -312,7 +316,7 @@ def test_reduction_add_inner(self):
         ir,
         "reduction_add_inner",
         4,
-        [InputSpec((1024, 32), 1), InputSpec((1,), 0)],
+        [InputSpec((1024, 32)), InputSpec((1,))],
         (1024,),
         np.int32,
         lambda input, init: np.sum(input, axis=1) + init,
@@ -348,7 +352,7 @@ def test_reduction_add_outer(self):
         ir,
         "reduction_add_outer",
         4,
-        [InputSpec((1024, 32), 1), InputSpec((1,), 0)],
+        [InputSpec((1024, 32)), InputSpec((1,))],
         (32,),
         np.float32,
         lambda input, init: np.sum(input, axis=0),
@@ -381,7 +385,7 @@ def test_reduction_middle(self):
         ir,
         "reduction_add_middle",
         1,
-        [InputSpec((8, 4, 2), 1), InputSpec((1,), 0)],
+        [InputSpec((8, 4, 2)), InputSpec((1,))],
         (8, 2),
         np.float32,
         lambda input, init: np.sum(input, axis=1),
@@ -414,7 +418,7 @@ def test_reduction_outer_inner(self):
         ir,
         "reduction_add_outer_inner",
         1,
-        [InputSpec((8, 4, 2), 1), InputSpec((1,), 0)],
+        [InputSpec((8, 4, 2)), InputSpec((1,))],
         (4,),
         np.float32,
         lambda input, init: np.sum(input, axis=(0, 2)),
@@ -439,7 +443,7 @@ def test_broadcast_in_dim_inner(self):
         ir,
         "broadcast_in_dim_inner",
         1,
-        [InputSpec((4,), None)],
+        [InputSpec((4,))],
         (32, 4),
         np.float32,
         lambda input: np.broadcast_to(input, (32, 4)),
@@ -464,7 +468,7 @@ def test_broadcast_in_dim_outer(self):
         ir,
         "broadcast_in_dim_outer",
         1,
-        [InputSpec((4,), None)],
+        [InputSpec((4,))],
         (4, 32),
         np.float32,
         lambda input: np.transpose(np.broadcast_to(input, (32, 4))),