Added 'expand_dims' to transform external calls

apache · Nov 13, 2023 · c30fbe8 · c30fbe8
1 parent 2b598a9
commit c30fbe8
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/rust/tvm/src/ir/relay/transform.rs b/rust/tvm/src/ir/relay/transform.rs
@@ -19,4 +19,7 @@ external! {
 
     #[name("relay.op._make.strided_slice")]
     pub fn strided_slice(data: Expr, begin: Array<PrimExpr>, end: Array<PrimExpr>, strides: Array<PrimExpr>, slice_mode: TVMString, axes: Array<PrimExpr>) -> Expr;
+
+    #[name("relay.op._make.expand_dims")]
+    pub fn expand_dims(data: Expr, axis: i32, num_newaxis: i32) -> Expr;
 }