abergeron commented on a change in pull request #4303: [TOPI][Relay][OP] Add a
strided_set operation.
URL: https://github.com/apache/incubator-tvm/pull/4303#discussion_r346912399
##########
File path: topi/python/topi/transform.py
##########
@@ -155,6 +157,78 @@ def strided_slice(a, begin, end, strides=None):
strides = []
return cpp.strided_slice(a, begin, end, strides)
[email protected]_scope(tag=tag.INJECTIVE+",strided_set")
+def strided_set(a, v, begin, end, strides=None):
+ """Set slice of an array.
+
+ Parameters
+ ----------
+ a : tvm.Tensor
+ The tensor to be sliced.
+
+ v : tvm.Tensor
+ The values to set
+
+ begin: list of Expr
+ The indices to begin with in the slicing.
+
+ end: list of Expr
+ Indicies indicating end of the slice.
+
+ strides: list of Expr, optional
+ Specifies the stride values, it can be negative
+ in that case, the input tensor will be reversed
+ in that particular axis.
+
+ Returns
+ -------
+ ret : tvm.Tensor
+ """
+ n = len(a.shape)
+ if strides is None:
+ strides = [1] * n
+
+ if len(begin) != n:
+ raise ValueError("size mismatch")
+ if len(end) != n:
+ raise ValueError("size mismatch")
+ if len(strides) != n:
+ raise ValueError("size mismatch")
+
+ begin = list(map(tvm.convert, begin))
+ end = list(map(tvm.convert, end))
+ strides = list(map(tvm.convert, strides))
+
+ def _max(a, b):
+ return tvm.expr.Select(a > b, a, b)
+
+ # Convert negative indexes
+ for i in range(n):
+ begin[i] = tvm.if_then_else(begin[i] < 0,
+ begin[i] + a.shape[i],
Review comment:
By OOB do you mean Out of Bounds of the indexed array? If yes, that is not
a problem because the code will not try to fetch indices that are out the array
shape in any cases.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
