Fix unsupported dt for Metal Cast

sonos · Oct 1, 2024 · 8441bf3 · 8441bf3
1 parent 9a53232
commit 8441bf3
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/metal/src/ops/cast.rs b/metal/src/ops/cast.rs
@@ -1,5 +1,5 @@
+use crate::kernels;
 use crate::tensor::MetalTensorExt;
-use crate::{kernels, MetalTensor};
 use tract_core::internal::*;
 
 #[derive(Debug, Clone, Hash, PartialEq, Eq)]
@@ -8,8 +8,12 @@ pub struct MetalCast {
 }
 
 impl MetalCast {
+    pub fn is_supported_dt(dt: DatumType) -> bool {
+        kernels::array::Cast::is_supported_dt(dt)
+    }
+
     pub fn new(to: DatumType) -> Option<Self> {
-        MetalTensor::is_supported_dt(to).then(|| Self { to })
+        Self::is_supported_dt(to).then_some(Self { to })
     }
 }
 

diff --git a/metal/src/transform.rs b/metal/src/transform.rs
@@ -180,7 +180,10 @@ impl Translate<TypedFact, Box<dyn TypedOp>, TypedFact, Box<dyn TypedOp>> for Met
         } else if let Some(op) = node.op_as::<Const>() {
             ops::MetalConst::new(op.0.clone())?.map(|o| -> Box<dyn TypedOp> { Box::new(o) })
         } else if let Some(op) = node.op_as::<Cast>() {
-            ops::MetalCast::new(op.to).map(|o| -> Box<dyn TypedOp> { Box::new(o) })
+            check_in_dts_are_supported(source, node.id, ops::MetalCast::is_supported_dt)?
+                .then(|| ops::MetalCast::new(op.to))
+                .flatten()
+                .map(|o| -> Box<dyn TypedOp> { Box::new(o) })
         } else if let Some(op) = node.op_as::<AxisOp>() {
             ops::MetalAxisOp::from_tract_core(op.clone())
                 .map(|o| -> Box<dyn TypedOp> { Box::new(o) })