Merge pull request #648 from robertknight/shape-start-end-attrs

Support `start` and `end` attributes for `Shape` operator

Author: robertknight

Diff for: rten-convert/rten_convert/converter.py

@@ -991,8 +991,13 @@ def op_node_from_onnx_operator(
             )
 
         case "Shape":
-            op_reader.check_attr("end", "int", 0)
-            op_reader.check_attr("start", "int", 0)
+            attrs = sg.ShapeAttrsT()
+            start = op_reader.get_attr("start", "int", None)
+            if start is not None:
+                attrs.start = start
+            end = op_reader.get_attr("end", "int", None)
+            if end is not None:
+                attrs.end = end
 
         case "Softmax":
             attrs = sg.SoftmaxAttrsT()

Diff for: rten-convert/rten_convert/schema_generated.py

@@ -204,6 +204,7 @@ class OperatorAttrs(object):
     QuantizeLinearAttrs = 42
     DepthToSpaceAttrs = 43
     CastLikeAttrs = 44
+    ShapeAttrs = 45
 
 def OperatorAttrsCreator(unionType, table):
     from flatbuffers.table import Table
@@ -297,6 +298,8 @@ def OperatorAttrsCreator(unionType, table):
         return DepthToSpaceAttrsT.InitFromBuf(table.Bytes, table.Pos)
     if unionType == OperatorAttrs.CastLikeAttrs:
         return CastLikeAttrsT.InitFromBuf(table.Bytes, table.Pos)
+    if unionType == OperatorAttrs.ShapeAttrs:
+        return ShapeAttrsT.InitFromBuf(table.Bytes, table.Pos)
     return None
 
 
@@ -4631,6 +4634,96 @@ def Pack(self, builder):
         return scatterNdattrs
 
 
+class ShapeAttrs(object):
+    __slots__ = ['_tab']
+
+    @classmethod
+    def GetRootAs(cls, buf, offset=0):
+        n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, offset)
+        x = ShapeAttrs()
+        x.Init(buf, n + offset)
+        return x
+
+    @classmethod
+    def GetRootAsShapeAttrs(cls, buf, offset=0):
+        """This method is deprecated. Please switch to GetRootAs."""
+        return cls.GetRootAs(buf, offset)
+    @classmethod
+    def ShapeAttrsBufferHasIdentifier(cls, buf, offset, size_prefixed=False):
+        return flatbuffers.util.BufferHasIdentifier(buf, offset, b"\x52\x54\x45\x4E", size_prefixed=size_prefixed)
+
+    # ShapeAttrs
+    def Init(self, buf, pos):
+        self._tab = flatbuffers.table.Table(buf, pos)
+
+    # ShapeAttrs
+    def Start(self):
+        o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
+        if o != 0:
+            return self._tab.Get(flatbuffers.number_types.Int32Flags, o + self._tab.Pos)
+        return None
+
+    # ShapeAttrs
+    def End(self):
+        o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(6))
+        if o != 0:
+            return self._tab.Get(flatbuffers.number_types.Int32Flags, o + self._tab.Pos)
+        return None
+
+def ShapeAttrsStart(builder):
+    builder.StartObject(2)
+
+def ShapeAttrsAddStart(builder, start):
+    builder.PrependInt32Slot(0, start, None)
+
+def ShapeAttrsAddEnd(builder, end):
+    builder.PrependInt32Slot(1, end, None)
+
+def ShapeAttrsEnd(builder):
+    return builder.EndObject()
+
+
+
+class ShapeAttrsT(object):
+
+    # ShapeAttrsT
+    def __init__(self):
+        self.start = None  # type: Optional[int]
+        self.end = None  # type: Optional[int]
+
+    @classmethod
+    def InitFromBuf(cls, buf, pos):
+        shapeAttrs = ShapeAttrs()
+        shapeAttrs.Init(buf, pos)
+        return cls.InitFromObj(shapeAttrs)
+
+    @classmethod
+    def InitFromPackedBuf(cls, buf, pos=0):
+        n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, pos)
+        return cls.InitFromBuf(buf, pos+n)
+
+    @classmethod
+    def InitFromObj(cls, shapeAttrs):
+        x = ShapeAttrsT()
+        x._UnPack(shapeAttrs)
+        return x
+
+    # ShapeAttrsT
+    def _UnPack(self, shapeAttrs):
+        if shapeAttrs is None:
+            return
+        self.start = shapeAttrs.Start()
+        self.end = shapeAttrs.End()
+
+    # ShapeAttrsT
+    def Pack(self, builder):
+        ShapeAttrsStart(builder)
+        ShapeAttrsAddStart(builder, self.start)
+        ShapeAttrsAddEnd(builder, self.end)
+        shapeAttrs = ShapeAttrsEnd(builder)
+        return shapeAttrs
+
+
 class SoftmaxAttrs(object):
     __slots__ = ['_tab']
 
@@ -5223,7 +5316,7 @@ class OperatorNodeT(object):
     def __init__(self):
         self.type = 0  # type: int
         self.attrsType = 0  # type: int
-        self.attrs = None  # type: Union[None, ArgMaxAttrsT, AveragePoolAttrsT, BatchNormalizationAttrsT, CastAttrsT, ConcatAttrsT, ConstantOfShapeAttrsT, ConvAttrsT, ConvTransposeAttrsT, FlattenAttrsT, GatherAttrsT, GemmAttrsT, GRUAttrsT, LeakyReluAttrsT, LSTMAttrsT, MaxPoolAttrsT, ReduceMeanAttrsT, ReshapeAttrsT, ResizeAttrsT, SplitAttrsT, SoftmaxAttrsT, TransposeAttrsT, ModAttrsT, ScatterElementsAttrsT, OneHotAttrsT, TopKAttrsT, HardSigmoidAttrsT, TriluAttrsT, ScatterNDAttrsT, NonMaxSuppressionAttrsT, LayerNormalizationAttrsT, RandomUniformAttrsT, EluAttrsT, RandomUniformLikeAttrsT, RandomNormalAttrsT, RandomNormalLikeAttrsT, GatherNDAttrsT, GeluAttrsT, EinsumAttrsT, IfAttrsT, PadAttrsT, DequantizeLinearAttrsT, QuantizeLinearAttrsT, DepthToSpaceAttrsT, CastLikeAttrsT]
+        self.attrs = None  # type: Union[None, ArgMaxAttrsT, AveragePoolAttrsT, BatchNormalizationAttrsT, CastAttrsT, ConcatAttrsT, ConstantOfShapeAttrsT, ConvAttrsT, ConvTransposeAttrsT, FlattenAttrsT, GatherAttrsT, GemmAttrsT, GRUAttrsT, LeakyReluAttrsT, LSTMAttrsT, MaxPoolAttrsT, ReduceMeanAttrsT, ReshapeAttrsT, ResizeAttrsT, SplitAttrsT, SoftmaxAttrsT, TransposeAttrsT, ModAttrsT, ScatterElementsAttrsT, OneHotAttrsT, TopKAttrsT, HardSigmoidAttrsT, TriluAttrsT, ScatterNDAttrsT, NonMaxSuppressionAttrsT, LayerNormalizationAttrsT, RandomUniformAttrsT, EluAttrsT, RandomUniformLikeAttrsT, RandomNormalAttrsT, RandomNormalLikeAttrsT, GatherNDAttrsT, GeluAttrsT, EinsumAttrsT, IfAttrsT, PadAttrsT, DequantizeLinearAttrsT, QuantizeLinearAttrsT, DepthToSpaceAttrsT, CastLikeAttrsT, ShapeAttrsT]
         self.inputs = None  # type: List[int]
         self.outputs = None  # type: List[int]
 

Diff for: src/graph.rs

@@ -1369,7 +1369,7 @@ mod tests {
         let input_b_id = g.add_value(Some("input_b"), None, None);
 
         let (add_op, add_out) = g.add_simple_op("add", Add {}, &[input_a_id, input_b_id]);
-        let (shape_op, shape_out) = g.add_simple_op("shape", Shape {}, &[input_a_id]);
+        let (shape_op, shape_out) = g.add_simple_op("shape", Shape::default(), &[input_a_id]);
 
         // The execution plan could run operators in either order and produce
         // the correct output. Since the `Add` op has the _potential_ to run in
@@ -1624,7 +1624,12 @@ mod tests {
         // as opposed to passing a shorter input list. This enables omitting
         // an input but still providing subsequent ones.
         let output = g.add_value(None, None, None);
-        g.add_op(Some("shape"), Box::new(Shape {}), &[None], &[Some(output)]);
+        g.add_op(
+            Some("shape"),
+            Box::new(Shape::default()),
+            &[None],
+            &[Some(output)],
+        );
 
         let results = g.run(vec![], &[output], None, None);
 

Diff for: src/model.rs

@@ -887,7 +887,7 @@ mod tests {
     use crate::ops;
     use crate::ops::{
         BoxOrder, CoordTransformMode, DataType, DepthToSpaceMode, NearestMode, OpError, Output,
-        ResizeMode, Scalar,
+        ResizeMode, Scalar, Shape,
     };
     use crate::{ModelLoadError, OpRegistry, ReadOpError};
 
@@ -1214,7 +1214,12 @@ mod tests {
 
         let output_node = graph_builder.add_value("output", None, None);
         graph_builder.add_output(output_node);
-        graph_builder.add_operator("shape", OpType::Shape, &[None], &[output_node]);
+        graph_builder.add_operator(
+            "shape",
+            OpType::Shape(Shape::default()),
+            &[None],
+            &[output_node],
+        );
 
         let graph = graph_builder.finish();
         builder.set_graph(graph);
@@ -1597,7 +1602,10 @@ mod tests {
 
         add_operator!(Round, [input_node]);
 
-        add_operator!(Shape, [input_node]);
+        add_operator!(Shape, [input_node], {
+            start: Some(1),
+            end: Some(-1),
+        });
         add_operator!(Sigmoid, [input_node]);
         add_operator!(Sign, [input_node]);
         add_operator!(Sin, [input_node]);

Diff for: src/model_builder.rs

@@ -12,7 +12,7 @@ use crate::ops::{
     Gelu, Gemm, HardSigmoid, InstanceNormalization, LayerNormalization, LeakyRelu, LogSoftmax,
     MaxPool, Mod, NearestMode, NonMaxSuppression, OneHot, Padding, QuantizeLinear, ReduceMax,
     ReduceMean, ReduceMin, ReduceProd, ReduceSum, ReduceSumSquare, Reshape, Resize, ResizeMode,
-    Scalar, ScatterElements, ScatterReduction, Softmax, Split, TopK, Transpose, Trilu,
+    Scalar, ScatterElements, ScatterReduction, Shape, Softmax, Split, TopK, Transpose, Trilu,
 };
 use crate::schema_generated as sg;
 
@@ -119,7 +119,7 @@ pub enum OpType<'a> {
     Round,
     QuantizeLinear(QuantizeLinear),
     ScatterElements(ScatterElements),
-    Shape,
+    Shape(Shape),
     Sigmoid,
     Sign,
     Sin,
@@ -828,7 +828,12 @@ impl<'mb, 'a> GraphBuilder<'mb, 'a> {
                     }
                 })
             }
-            OpType::Shape => op!(Shape),
+            OpType::Shape(args) => op_with_attrs!(Shape, ShapeAttrs, {
+                sg::ShapeAttrsArgs {
+                    start: args.start,
+                    end: args.end,
+                }
+            }),
             OpType::Sigmoid => op!(Sigmoid),
             OpType::Slice => op!(Slice),
             OpType::Sin => op!(Sin),

Diff for: src/op_registry.rs

@@ -847,7 +847,21 @@ impl_read_op!(
         })
     }
 );
-impl_read_op!(Shape);
+
+impl ReadOp for ops::Shape {
+    fn op_type() -> sg::OperatorType {
+        OperatorType::Shape
+    }
+
+    fn read(op: &OperatorNode, _ctx: &dyn OpLoadContext) -> Result<Self, ReadOpError> {
+        // Shape attributes are optional for backwards compatibility
+        let attrs = op.attrs_as_shape_attrs();
+        let start = attrs.and_then(|a| a.start());
+        let end = attrs.and_then(|a| a.end());
+        Ok(ops::Shape { start, end })
+    }
+}
+
 impl_read_op!(Sigmoid);
 impl_read_op!(Sign);
 impl_read_op!(Sin);