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_r347155491
##########
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:
Even if the logical result is very similar to strided_slice the
implementation has a different.
In this operation: `a[begin:end:stride] = b`
The core kernel loops over all valid indexes for `a` and check if that index
is part of the values selected by the combination of `begin`, `end` and
`stride`. If it is, it will compute the corresponding index in `b` and map the
output to that value. Otherwise it will pick up the value from `a` at that
index.
In all cases it doesn't matter if `begin`, `end`, or `stride` doesn't fall
within the bounds of `a` because they are never used to directly or indirectly
index into `a`.
----------------------------------------------------------------
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
