Merge pull request #646 from robertknight/cast-like-op

Implement `CastLike` operator

Author: robertknight

rten-convert/rten_convert/schema_generated.py

@@ -118,6 +118,7 @@ class OperatorType(object):
     MatMulInteger = 108
     DepthToSpace = 109
     ConvInteger = 110
+    CastLike = 111
 
 
 class RNNDirection(object):
@@ -201,6 +202,7 @@ class OperatorAttrs(object):
     DequantizeLinearAttrs = 41
     QuantizeLinearAttrs = 42
     DepthToSpaceAttrs = 43
+    CastLikeAttrs = 44
 
 def OperatorAttrsCreator(unionType, table):
     from flatbuffers.table import Table
@@ -292,6 +294,8 @@ def OperatorAttrsCreator(unionType, table):
         return QuantizeLinearAttrsT.InitFromBuf(table.Bytes, table.Pos)
     if unionType == OperatorAttrs.DepthToSpaceAttrs:
         return DepthToSpaceAttrsT.InitFromBuf(table.Bytes, table.Pos)
+    if unionType == OperatorAttrs.CastLikeAttrs:
+        return CastLikeAttrsT.InitFromBuf(table.Bytes, table.Pos)
     return None
 
 
@@ -873,6 +877,71 @@ def Pack(self, builder):
         return castAttrs
 
 
+class CastLikeAttrs(object):
+    __slots__ = ['_tab']
+
+    @classmethod
+    def GetRootAs(cls, buf, offset=0):
+        n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, offset)
+        x = CastLikeAttrs()
+        x.Init(buf, n + offset)
+        return x
+
+    @classmethod
+    def GetRootAsCastLikeAttrs(cls, buf, offset=0):
+        """This method is deprecated. Please switch to GetRootAs."""
+        return cls.GetRootAs(buf, offset)
+    @classmethod
+    def CastLikeAttrsBufferHasIdentifier(cls, buf, offset, size_prefixed=False):
+        return flatbuffers.util.BufferHasIdentifier(buf, offset, b"\x52\x54\x45\x4E", size_prefixed=size_prefixed)
+
+    # CastLikeAttrs
+    def Init(self, buf, pos):
+        self._tab = flatbuffers.table.Table(buf, pos)
+
+def CastLikeAttrsStart(builder):
+    builder.StartObject(0)
+
+def CastLikeAttrsEnd(builder):
+    return builder.EndObject()
+
+
+
+class CastLikeAttrsT(object):
+
+    # CastLikeAttrsT
+    def __init__(self):
+        pass
+
+    @classmethod
+    def InitFromBuf(cls, buf, pos):
+        castLikeAttrs = CastLikeAttrs()
+        castLikeAttrs.Init(buf, pos)
+        return cls.InitFromObj(castLikeAttrs)
+
+    @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, castLikeAttrs):
+        x = CastLikeAttrsT()
+        x._UnPack(castLikeAttrs)
+        return x
+
+    # CastLikeAttrsT
+    def _UnPack(self, castLikeAttrs):
+        if castLikeAttrs is None:
+            return
+
+    # CastLikeAttrsT
+    def Pack(self, builder):
+        CastLikeAttrsStart(builder)
+        castLikeAttrs = CastLikeAttrsEnd(builder)
+        return castLikeAttrs
+
+
 class ConcatAttrs(object):
     __slots__ = ['_tab']
 
@@ -5153,7 +5222,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]
+        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.inputs = None  # type: List[int]
         self.outputs = None  # type: List[int]
 

src/model.rs

@@ -1330,6 +1330,7 @@ mod tests {
         );
 
         add_operator!(Cast, [input_node], { to: ops::DataType::Float });
+        add_operator!(CastLike, [input_node, input_node], {});
         add_operator!(Ceil, [input_node]);
 
         let clip_min = graph_builder.add_constant(Tensor::from(1.).view());

src/model_builder.rs

@@ -6,13 +6,13 @@ use crate::graph::{Dimension, NodeId};
 use crate::header::Header;
 use crate::number::LeBytes;
 use crate::ops::{
-    ArgMax, ArgMin, AveragePool, BatchNormalization, BoxOrder, Cast, Concat, ConstantOfShape, Conv,
-    ConvInteger, ConvTranspose, CoordTransformMode, DataType, DepthToSpace, DepthToSpaceMode,
-    DequantizeLinear, Einsum, Elu, Flatten, Gather, GatherElements, GatherND, 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,
+    ArgMax, ArgMin, AveragePool, BatchNormalization, BoxOrder, Cast, CastLike, Concat,
+    ConstantOfShape, Conv, ConvInteger, ConvTranspose, CoordTransformMode, DataType, DepthToSpace,
+    DepthToSpaceMode, DequantizeLinear, Einsum, Elu, Flatten, Gather, GatherElements, GatherND,
+    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,
 };
 use crate::schema_generated as sg;
 
@@ -39,6 +39,7 @@ pub enum OpType<'a> {
     AveragePool(AveragePool),
     BatchNormalization(BatchNormalization),
     Cast(Cast),
+    CastLike(CastLike),
     Ceil,
     Clip,
     Concat(Concat),
@@ -449,6 +450,9 @@ impl<'mb, 'a> GraphBuilder<'mb, 'a> {
                     to: convert_dtype(args.to),
                 }
             ),
+            OpType::CastLike(_args) => {
+                op_with_attrs!(CastLike, CastLikeAttrs, sg::CastLikeAttrsArgs {})
+            }
             OpType::Ceil => op!(Ceil),
             OpType::Clip => op!(Clip),
             OpType::Concat(args) => op_with_attrs!(

src/op_registry.rs

@@ -92,6 +92,7 @@ impl OpRegistry {
         register_op!(AveragePool);
         register_op!(BatchNormalization);
         register_op!(Cast);
+        register_op!(CastLike);
         register_op!(Ceil);
         register_op!(Clip);
         register_op!(Concat);
@@ -438,6 +439,11 @@ impl_read_op!(Cast, attrs_as_cast_attrs, |attrs: sg::CastAttrs| {
     let to = convert_dtype(attrs.to())?;
     Ok(ops::Cast { to })
 });
+impl_read_op!(
+    CastLike,
+    attrs_as_cast_like_attrs,
+    |_attrs: sg::CastLikeAttrs| { Ok(ops::CastLike {}) }
+);
 impl_read_op!(Ceil);
 impl_read_op!(Clip);
 impl_read_op!(Concat, attrs_as_concat_attrs, axis);

src/ops/convert.rs

@@ -79,13 +79,49 @@ impl Operator for Cast {
     }
 }
 
+#[derive(Debug)]
+pub struct CastLike {}
+
+impl Operator for CastLike {
+    fn name(&self) -> &str {
+        "CastLike"
+    }
+
+    fn run(&self, pool: &TensorPool, inputs: InputList) -> Result<OutputList, OpError> {
+        let input = inputs.require(0)?;
+        let to_type = inputs.require(1)?.dtype();
+        cast(pool, input, to_type).into_op_result()
+    }
+
+    fn can_run_in_place(&self) -> bool {
+        true
+    }
+
+    fn run_in_place(
+        &self,
+        pool: &TensorPool,
+        input: Output,
+        other: InputList,
+    ) -> Result<Output, OpError> {
+        let to_type = other.require(0)?.dtype();
+
+        if input.dtype() == to_type {
+            Ok(input)
+        } else {
+            let converted = cast(pool, input.as_input(), to_type)?;
+            input.add_to_pool(pool);
+            Ok(converted)
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use rten_tensor::Tensor;
     use rten_testing::TestCases;
 
     use crate::ops::tests::new_pool;
-    use crate::ops::{Cast, DataType, Operator, Output};
+    use crate::ops::{Cast, CastLike, DataType, Operator, Output};
 
     #[test]
     fn test_cast() {
@@ -159,4 +195,36 @@ mod tests {
             assert_eq!(result, case.expected);
         })
     }
+
+    #[test]
+    fn test_cast_like() {
+        #[derive(Debug)]
+        struct Case {
+            input: Output,
+            other: Output,
+            expected: Output,
+        }
+
+        // `CastLike` uses the same conversions as the `Cast` operator,
+        // so these tests don't check all data type combinations, only that the
+        // target type is taken from the second argument.
+        let cases = [
+            // i32 -> f32
+            Case {
+                input: Tensor::from([0i32, 1, 2]).into(),
+                other: Tensor::from([0f32]).into(),
+                expected: Tensor::from([0., 1., 2.]).into(),
+            },
+        ];
+
+        cases.test_each(|case| {
+            let pool = new_pool();
+            let cast_op = CastLike {};
+            let result = cast_op
+                .run(&pool, (&case.input, &case.other).into())
+                .unwrap()
+                .remove(0);
+            assert_eq!(result, case.expected);
+        })
+    }
 }

src/ops/mod.rs

@@ -72,7 +72,7 @@ pub use binary_elementwise::{
 pub use concat::{concat, tile, Concat, Tile};
 pub use control_flow::If;
 pub use conv::{conv, conv_integer, conv_transpose, Conv, ConvInteger, ConvTranspose};
-pub use convert::Cast;
+pub use convert::{Cast, CastLike};
 pub use einsum::{einsum, Einsum};
 pub use gather::{
     gather, gather_elements, gather_nd, scatter_elements, scatter_nd, Gather, GatherElements,
@@ -255,6 +255,16 @@ pub enum Input<'a> {
 }
 
 impl Input<'_> {
+    /// Return the data type of elements in this tensor.
+    pub fn dtype(&self) -> DataType {
+        match self {
+            Self::FloatTensor(_) => DataType::Float,
+            Self::Int32Tensor(_) => DataType::Int32,
+            Self::Int8Tensor(_) => DataType::Int8,
+            Self::UInt8Tensor(_) => DataType::UInt8,
+        }
+    }
+
     pub fn to_output(&self) -> Output {
         match self {
             Input::FloatTensor(t) => t.to_tensor().into(),

src/schema.fbs

@@ -124,6 +124,7 @@ enum OperatorType: ubyte {
   MatMulInteger,
   DepthToSpace,
   ConvInteger,
+  CastLike,
 }
 
 enum RNNDirection: ubyte {
@@ -218,6 +219,7 @@ union OperatorAttrs {
   DequantizeLinearAttrs,
   QuantizeLinearAttrs,
   DepthToSpaceAttrs,
+  CastLikeAttrs,
 }
 
 table ArgMaxAttrs {
@@ -245,6 +247,8 @@ table CastAttrs {
   to:DataType;
 }
 
+table CastLikeAttrs {}
+
 table ConcatAttrs {
   axis:int;
 }