chore: tidy up

huggingface · Oct 29, 2024 · 5a018ce · 5a018ce
1 parent be77442
commit 5a018ce
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 6 deletions.
diff --git a/crates/ratchet-core/src/ops/rope.rs b/crates/ratchet-core/src/ops/rope.rs
@@ -279,8 +279,6 @@ import mlx.core as mx
 import mlx.nn as nn
 import numpy as np
 
-mx.set_default_device(mx.cpu)
-
 def mlx_rope(input, dim, offset):
     rope = nn.RoPE(dim)
     mx_input = mx.array(input)
@@ -332,7 +330,7 @@ def mlx_rope(input, dim, offset):
         offset: usize,
     }
 
-    #[proptest(cases = 8)]
+    #[proptest(cases = 16)]
     fn test_rope_gpu(prob: RoPEProblem) {
         let RoPEProblem {
             BS,
@@ -359,9 +357,8 @@ def mlx_rope(input, dim, offset):
             SL,
             HD,
             dim,
-            mut offset,
+            offset,
         } = prob;
-        offset = 0;
         println!(
             "BS = {}, NH = {}, SL = {}, HD = {}, rope_dim = {}, offset = {}",
             BS, NH, SL, HD, dim, offset

diff --git a/crates/ratchet-core/src/storage/cpu_buffer.rs b/crates/ratchet-core/src/storage/cpu_buffer.rs
@@ -89,7 +89,7 @@ impl CPUBuffer {
     }
 
     pub fn from_slice<T: NoUninit>(data: &[T], shape: &Shape) -> Self {
-        //assert_eq!(data.len(), shape.numel());
+        assert_eq!(data.len(), shape.numel());
         let bytes: &[u8] = bytemuck::cast_slice(data);
         Self::from_bytes(bytes, std::mem::align_of::<T>())
     }