chore: import upstream snapshot with attribution

python/tvm/topi/testing/strided_slice_python.py

@@ -0,0 +1,129 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""strided_slice/set in python"""
+
+
+def strided_slice_python(data, begin, end, strides, slice_mode="end", axes=None):
+    """Python version of strided slice operator.
+
+    Parameters
+    ----------
+    data : numpy.ndarray
+        Input data
+
+    begin : list
+        Beginning of the slices.
+
+    end : list
+        End of the slices.
+
+    strides : list
+        The stride of each slice.
+
+    slice_mode : str, optional
+        The slice mode [end, size].
+
+        - ``"end"``: The default slice mode, ending indices for the slice.
+        - ``"size"``: The input strides will be ignored, input end in this mode indicates
+          the size of a slice starting at the location specified by begin. If end[i] is -1,
+          all remaining elements in that dimension are included in the slice.
+
+    axes : list, optional
+        Axes along which slicing is applied
+
+    Returns
+    -------
+    result : numpy.ndarray
+        The sliced result.
+    """
+    strides = [] if strides is None else strides
+    if axes is not None:
+        rank = len(data.shape)
+        new_begin = [0] * rank
+        new_end = [data.shape[i] for i in range(rank)]
+        new_strides = [1] * rank
+
+        for i, axis in enumerate(axes):
+            new_begin[axis] = begin[i]
+            new_end[axis] = end[i]
+            if len(strides) > i:
+                new_strides[axis] = strides[i]
+
+        begin = new_begin
+        end = new_end
+        strides = new_strides
+
+    slices = []
+    for i in range(len(data.shape)):
+        new_stride = None
+        if slice_mode == "end" and i < len(strides):
+            new_stride = strides[i]
+
+        new_begin = begin[i] if i < len(begin) else None
+        if i >= len(end):
+            new_end = None
+        elif slice_mode == "size":
+            if end[i] < 0:
+                new_end = None
+            else:
+                new_end = new_begin + end[i]
+        else:
+            new_end = end[i]
+
+        slices.append(slice(new_begin, new_end, new_stride))
+
+    return data[tuple(slices)]
+
+
+def strided_set_python(data, v, begin, end, strides):
+    """Python version of strided slice operator.
+
+    Parameters
+    ----------
+    data : numpy.ndarray
+        Input data
+
+    v : numpy.ndarray
+        Value data
+
+    begin : list
+        Beginning of the slices.
+
+    end : list
+        End of the slices.
+
+    strides : list
+        The stride of each slice.
+
+    Returns
+    -------
+    result : numpy.ndarray
+        The updated result.
+    """
+    strides = [] if strides is None else strides
+    slices = []
+    res = data.copy()
+    for i in range(len(data.shape)):
+        slices.append(
+            slice(
+                begin[i] if i < len(begin) else None,
+                end[i] if i < len(end) else None,
+                strides[i] if i < len(strides) else None,
+            )
+        )
+    res[tuple(slices)] = v
+    return res