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 never try to fetch indices that are out the
array shape.
This code just tries to handle numpy-style negative indexing (starts from
the end of the array) just like strided_slice does.
----------------------------------------------------------------
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
