AddOp(aten_replication_pad1d) | feat(torchlib) (#1203)

- A large change in extra_opinfo.py, is just because re-sort the functions by name.
microsoft · Dec 5, 2023 · e3a9366 · e3a9366
1 parent 173cd61
commit e3a9366
Show file tree

Hide file tree

Showing 3 changed files with 1,117 additions and 1,079 deletions.
diff --git a/onnxscript/function_libs/torch_lib/ops/nn.py b/onnxscript/function_libs/torch_lib/ops/nn.py
@@ -1496,10 +1496,18 @@ def aten_relu6(self: TReal) -> TReal:
     return op.Min(op.Relu(self), six)
 
 
+@torch_op("aten::replication_pad1d")
 def aten_replication_pad1d(self: TensorType, padding: INT64) -> TensorType:
     """replication_pad1d(Tensor self, SymInt[2] padding) -> Tensor"""
 
-    raise NotImplementedError()
+    # assert len(padding) == 2
+    # Input of padding argument should be [x,y], need change to onnx format [0, x, 0, y]
+    start = op.Slice(padding, [0], [1], axes=[0])
+    end = op.Slice(padding, [1], [2], axes=[0])
+    padding_onnx = op.Concat(
+        op.Constant(value_ints=[0]), start, op.Constant(value_ints=[0]), end, axis=0
+    )
+    return op.Pad(self, padding_onnx, mode="edge")
 
 
 def aten_replication_pad1d_backward(