dont try to be smart with axes in block_quant einsum

sonos · Aug 23, 2024 · feb5aff · feb5aff
1 parent 99ffaa6
commit feb5aff
Show file tree

Hide file tree

Showing 5 changed files with 25 additions and 9 deletions.
diff --git a/core/src/ops/array/concat.rs b/core/src/ops/array/concat.rs
@@ -110,6 +110,8 @@ impl TypedOp for TypedConcat {
     fn slice(
         &self,
         patch: &mut TypedModelPatch,
+        _model: &TypedModel,
+        _node: &TypedNode,
         prefix: &str,
         inputs: &[OutletId],
         output_axis: usize,

diff --git a/core/src/ops/binary.rs b/core/src/ops/binary.rs
@@ -186,6 +186,8 @@ impl TypedOp for TypedBinOp {
     fn slice(
         &self,
         patch: &mut TypedModelPatch,
+        _model: &TypedModel,
+        _node: &TypedNode,
         prefix: &str,
         inputs: &[OutletId],
         _output_axis: usize,

diff --git a/core/src/ops/einsum/mod.rs b/core/src/ops/einsum/mod.rs
@@ -305,13 +305,19 @@ impl TypedOp for EinSum {
     fn slice(
         &self,
         patch: &mut TypedModelPatch,
+        model: &TypedModel,
+        node: &TypedNode,
         prefix: &str,
         inputs: &[OutletId],
         _output_axis: usize,
         _start: usize,
         _end: usize,
     ) -> TractResult<Option<TVec<OutletId>>> {
-        patch.wire_node(prefix, self.clone(), inputs).map(Some)
+        if model.node_input_facts(node.id)?.iter().any(|f| f.datum_type.is_opaque()) {
+            Ok(None)
+        } else {
+            patch.wire_node(prefix, self.clone(), inputs).map(Some)
+        }
     }
 
     #[allow(unused_variables)]

diff --git a/core/src/ops/mod.rs b/core/src/ops/mod.rs
@@ -216,6 +216,8 @@ pub trait TypedOp:
     fn slice(
         &self,
         patch: &mut TypedModelPatch,
+        model: &TypedModel,
+        node: &TypedNode,
         prefix: &str,
         inputs: &[OutletId],
         output_axis: usize,

diff --git a/core/src/optim/slice.rs b/core/src/optim/slice.rs
@@ -56,18 +56,22 @@ impl super::TypedPass for PushSliceUp {
                             }
                             wires.push(wire);
                         }
-                        let Some(wire) = node.op.slice(
-                            &mut patch,
-                            &format!(
-                                "{}.split-over-{}.{}..{}",
-                                &node.name, axis, start, end
-                                ),
+                        let Some(wire) = node
+                            .op
+                            .slice(
+                                &mut patch,
+                                model,
+                                node,
+                                &format!("{}.split-over-{}.{}..{}", &node.name, axis, start, end),
                                 &wires,
                                 axis,
                                 start,
                                 *end,
-                                ).with_context(|| format!("Calling slice on {node}"))? else {
-                            continue 'axis };
+                            )
+                            .with_context(|| format!("Calling slice on {node}"))?
+                        else {
+                            continue 'axis;
+                        };
                         splits.push(wire[0]);
                         start = *end;
                     }