small changes

jafioti · May 3, 2024 · 6d917dd · 6d917dd
1 parent 97b66a3
commit 6d917dd
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/crates/luminal_metal/src/tests/fp16.rs b/crates/luminal_metal/src/tests/fp16.rs
@@ -770,7 +770,7 @@ fn test_movement() {
     let data = random_vec(32);
     let mut cx = Graph::new();
     let a = cx.tensor::<R1<32>>().set(data.clone());
-    let b: GraphTensor<R1<42>> = a.pad(&[(0, 10)]).contiguous().retrieve();
+    let b: GraphTensor<R1<42>> = a.pad((0, 10)).contiguous().retrieve();
     let mut c: GraphTensor<R1<25>> = b
         .slice((..Expression::from(25),))
         .realize()

diff --git a/src/hl_ops/movement.rs b/src/hl_ops/movement.rs
@@ -145,12 +145,11 @@ impl<S: Shape> GraphTensor<S> {
         if n_dims > 1 {
             // View as single dimension of matrix with wider width
             let mat_size = (dim_size.big() + stride.big()) * number_of_windows.big();
-            let actual_size =
-                dim_size.big() * self.shape.dims[self.shape.indexes[n_dims - 1]].big();
+            let actual_size = dim_size * self.shape.dims[self.shape.indexes[n_dims - 1]];
             // Reshape into single dimension to pad
             self.shape.remove_dim(n_dims);
-            self.shape.dims[self.shape.indexes[n_dims - 1]] = actual_size.small();
-            self.shape.padding[self.shape.indexes[n_dims - 1]].1 = (mat_size - actual_size).small();
+            self.shape.dims[self.shape.indexes[n_dims - 1]] = actual_size;
+            self.shape.padding[self.shape.indexes[n_dims - 1]].1 = mat_size.small() - actual_size;
             self = self.contiguous();
             // Reshape back (mats should be full now)
             self.shape.add_dim(n_dims, dim_size + stride);