Add ops(as_strided) | feat(atenlib) (#545)

Co-authored-by: Justin Chu <[email protected]>

Author: xiaowuhu

onnxscript/function_libs/torch_aten/ops/core.py

@@ -503,12 +503,74 @@ def aten_argwhere(self: TensorType) -> TensorType:
     raise NotImplementedError()
 
 
+@torch_op("aten::as_strided", trace_only=True)
 def aten_as_strided(
-    self: TensorType, size: INT64, stride: INT64, storage_offset: Optional[INT64] = None
-) -> TensorType:
+    self: TTensor, size: INT64, stride: INT64, storage_offset: int = 0
+) -> TTensor:
     """as_strided(Tensor(a) self, SymInt[] size, SymInt[] stride, SymInt? storage_offset=None) -> Tensor(a)"""
 
-    raise NotImplementedError()
+    rank = len(stride)
+    return _aten_as_strided_onnx(self, size, stride, storage_offset, rank)
+
+
+@torch_op("aten::as_strided", private=True)
+def _aten_as_strided_onnx(
+    self: TTensor, size: INT64, stride: INT64, storage_offset: int = 0, rank: int = 0
+) -> TTensor:
+    # e.g. when size=[2,3,4], stride=[2,1,3], indices=[0]
+    # i = 0
+    # indices=[0], add_value=[0,3,6,9]
+    # expand(shape=[4]) to [0,0,0,0]
+    # then + add_value = [0,3,6,9]
+    # i = 1
+    # indices=[0,3,6,9], add_value=[0,1,2]
+    # expand(shape=[3,4] to [[0,3,6,9],[0,3,6,9],[0,3,6,9]]
+    # indices + add_value = [[0,3,6,9],[1,3,7,10],[2,5,8,11]]
+    # i = 2
+    # indices = [[0,3,6,9],[1,3,7,10],[2,5,8,11]], add_value=[0,2]
+    # expand(shape=[2,3,4]) to [[[0,3,6,9],[1,3,7,10],[2,5,8,11]]],[[0,3,6,9],[1,3,7,10],[2,5,8,11]]]
+    # indices + add_value = [[[0,3,6,9],[1,3,7,10],[2,5,8,11]]],[[2,5,8,11],[3,5,9,12],[4,7,10,13]]]
+    neg_1 = op.Constant(value_ints=[-1])
+    rank_tensor = op.Reshape(rank, neg_1)  # should be 3
+    # The final indices for op.Gather(data, indices), will be continually changed during the loop
+    indices = op.Constant(value_int=0)
+    one_seq = op.SequenceEmpty()
+    for i in range(rank):
+        # Get the index from back to front, should be 2,1,0 when to i=0,1,2
+        j = rank - i - 1
+        j_tensor = op.Reshape(j, neg_1)
+        # Get size according to index_j, should be 4,3,2 when i=0,1,2
+        size_dim_j = op.Gather(size, j_tensor, axis=0)
+        # Get right size according to index_j, should be [4],[3,4],[2,3,4] when i=0,1,2
+        size_after_j = op.Slice(size, j_tensor, rank_tensor)
+        # Get stride according to index_j, should be 3,1,2 when i=0,1,2
+        stride_dim_j = op.Gather(stride, j_tensor, axis=0)
+        indices = op.Expand(indices, size_after_j)
+        # When size[j]=4, stride[j]=3, then add_value = [0,1,2,3] * 3 = [0,3,6,9]
+        # When size[j]=3, stride[j]=1, then add_value = [0,1,2] * 1 = [0,1,2]
+        # When size[j]=2, stride[j]=2, then add_value = [0,1] * 2 = [0,2]
+        add_value = op.Range(0, size_dim_j, 1) * stride_dim_j
+        # Compute the shape for add_value for correct broadcasting
+        if i == 0:
+            # shape = [dim_size]
+            shape = size_dim_j
+        else:
+            # shape = [dim_size, 1, 1, ...], the count of 1 euqal to i
+            ones = op.ConcatFromSequence(one_seq, axis=0)
+            shape = op.Concat(op.Cast(size_dim_j, to=FLOAT.dtype), ones, axis=0)
+            shape = op.Cast(shape, to=INT64.dtype)
+
+        add_value = op.Reshape(add_value, shape)
+        # Broadcasting add value to indices according to size and stride value
+        indices = indices + add_value
+        # Dims after dim_size to reshape(add_value), should be [1],[1,1],[1,1,1] when i=0,1,2
+        one_seq = op.SequenceInsert(one_seq, op.Constant(value_floats=[1.0]))
+
+    self_flatten = op.Reshape(self, op.Constant(value_ints=[-1]))
+    indices = op.Add(indices, storage_offset)
+    result = op.Gather(self_flatten, indices)
+
+    return result
 
 
 def aten_as_strided_copy(

onnxscript/tests/function_libs/torch_aten/ops_correctness_test.py

@@ -675,6 +675,7 @@ def _where_input_wrangler(
     "arange": core_ops.aten_arange,
     "argmax": core_ops.aten_argmax,
     "argmin": core_ops.aten_argmin,
+    "as_strided": core_ops.aten_as_strided,
     "clamp": core_ops.aten_clamp,
     "col2im": nn_ops.aten_col2im,
     "cumsum": core_ops.aten_cumsum,
@@ -751,6 +752,11 @@ def _where_input_wrangler(
         reason="fixme: ORT shape inference error",
         test_class_name="TestOutputConsistencyFullGraph",
     ),
+    xfail(
+        "as_strided",
+        variant_name="partial_views",
+        reason="ONNX doesn't have partial view for tensor",
+    ),
     xfail(
         "chunk", reason="fixme: ORT error", test_class_name="TestOutputConsistencyFullGraph"
     ),