feature: cpu slice remaining dtypes

huggingface · Nov 4, 2024 · 08900c2 · 08900c2
1 parent 7aee6e2
commit 08900c2
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 6 deletions.
diff --git a/crates/ratchet-core/src/cpu/reindex.rs b/crates/ratchet-core/src/cpu/reindex.rs
@@ -1,6 +1,6 @@
-use crate::{CPUOperation, DType, OperationError, Reindex, Slice, Strides, Tensor, TensorDType};
-
 use super::utils::cpu_store_result;
+use crate::{CPUOperation, DType, OperationError, Reindex, Slice, Strides, Tensor, TensorDType};
+use half::{bf16, f16};
 
 impl CPUOperation for Reindex {
     fn apply_cpu(&self, dst: Tensor) -> Result<Tensor, OperationError> {
@@ -15,6 +15,10 @@ impl CPUOperation for Slice {
     fn apply_cpu(&self, dst: Tensor) -> Result<Tensor, OperationError> {
         match dst.dt() {
             DType::F32 => apply_slice::<f32>(self, dst),
+            DType::BF16 => apply_slice::<bf16>(self, dst),
+            DType::F16 => apply_slice::<f16>(self, dst),
+            DType::I32 => apply_slice::<i32>(self, dst),
+            DType::U32 => apply_slice::<u32>(self, dst),
             _ => todo!(),
         }
     }

diff --git a/crates/ratchet-core/src/ops/reindex/slice.rs b/crates/ratchet-core/src/ops/reindex/slice.rs
@@ -158,10 +158,9 @@ def slice(a):
         run_py_prg(prg.to_string(), &[a], &[], a.dt())
     }
 
-    fn run_reindex_trial(prob: SliceProblem) -> anyhow::Result<()> {
+    fn run_reindex_trial(prob: SliceProblem, device: Device) -> anyhow::Result<()> {
         let SliceProblem { op } = prob;
         println!("SLICE PROBLEM: {:?}", op);
-        let device = Device::request_device(DeviceRequest::CPU).unwrap();
         let a = op.src.clone();
 
         let a_gpu = a.to(&device)?;
@@ -173,8 +172,16 @@ def slice(a):
     }
 
     #[proptest(cases = 16)]
-    fn test_slice(prob: SliceProblem) {
+    fn test_slice_gpu(prob: SliceProblem) {
+        let _ = env_logger::builder().is_test(true).try_init();
+        let device = Device::request_device(DeviceRequest::GPU).unwrap();
+        run_reindex_trial(prob, device).unwrap();
+    }
+
+    #[proptest(cases = 16)]
+    fn test_slice_cpu(prob: SliceProblem) {
         let _ = env_logger::builder().is_test(true).try_init();
-        run_reindex_trial(prob).unwrap();
+        let device = Device::request_device(DeviceRequest::CPU).unwrap();
+        run_reindex_trial(prob, device).unwrap();
     }
 }