contiguous check

mimiquate · Jan 16, 2024 · 0a6ed8c · 0a6ed8c
1 parent 8c7f219
commit 0a6ed8c
Showing 1 changed file with 15 additions and 0 deletions.
diff --git a/native/candlex/src/ops.rs b/native/candlex/src/ops.rs
@@ -106,6 +106,10 @@ macro_rules! custom_unary_op {
                 use crate::metal_kernels;
                 use candle_core::backend::BackendStorage;
 
+                if !(layout.is_contiguous() && layout.stride()[layout.stride().len() - 1] == 1) {
+                    candle_core::bail!("Non contiguous not supported");
+                }
+
                 let device = storage.device();
                 let command_buffer = device.command_buffer()?;
                 let elem_count = layout.shape().elem_count();
@@ -222,6 +226,10 @@ macro_rules! custom_unary_bool_op {
                 use crate::metal_kernels;
                 use candle_core::{backend::BackendStorage, DType};
 
+                if !(layout.is_contiguous() && layout.stride()[layout.stride().len() - 1] == 1) {
+                    candle_core::bail!("Non contiguous not supported");
+                }
+
                 let device = storage.device();
                 let command_buffer = device.command_buffer()?;
                 let elem_count = layout.shape().elem_count();
@@ -463,6 +471,13 @@ macro_rules! custom_binary_bool_op {
                 use crate::metal_kernels;
                 use candle_core::{backend::BackendStorage, DType};
 
+                if !(l1.is_contiguous() && l1.stride()[l1.stride().len() - 1] == 1) {
+                    candle_core::bail!("Non contiguous not supported");
+                }
+                if !(l2.is_contiguous() && l2.stride()[l2.stride().len() - 1] == 1) {
+                    candle_core::bail!("Non contiguous not supported");
+                }
+
                 let device = s1.device();
                 let shape = l1.shape();
                 let elem_count = shape.elem_count();