Extend shape op to return a slice #3580
Conversation
Signed-off-by: Ganesan Ramalingam <grama@microsoft.com>
Signed-off-by: Ganesan Ramalingam <grama@microsoft.com>
|
ONNX-TensorRT can support either this or the original ( |
| auto* output_length = output_shape->add_dim(); | ||
|
|
||
|
|
||
| if (!hasNInputShapes(ctx, 1)) { | ||
| return; | ||
| } | ||
|
|
||
| if (ctx.getInputType(0)->tensor_type().has_shape()) { |
There was a problem hiding this comment.
is this if required? if (!hasNInputShapes(ctx, 1)) is checking the same thing right?
onnx/defs/tensor/defs.cc
Outdated
| int64_t end = getAttribute(ctx, "end", rank); | ||
| if (end < 0) | ||
| end += rank; | ||
| output_length->set_dim_value(end - start); |
There was a problem hiding this comment.
validate the end-start before setting dim_value?
make sure end -start is positive > 0
There was a problem hiding this comment.
@gramalingam @askhade : Is it worth explicitly calling out the behavior in the spec that by setting start > end (and start == end because of the exclusive nature of end) the resultant shape slice will be empty (over it being an error) ? Just by reading the spec, this portion was unclear to me until I saw the shape inference. (Sorry if I overlooked something obvious in the spec)
|
I think this approach is better over introducing Dim (PR #3574 ) |
Signed-off-by: Ganesan Ramalingam <grama@microsoft.com>
Signed-off-by: Ganesan Ramalingam <grama@microsoft.com>
Signed-off-by: Ganesan Ramalingam grama@microsoft.com
Description
Extend the shape op by adding optional attributes start/end so that the return value is input_tensor.shape[start:end].
Motivation and Context
Helps abbreviate a common pattern, eg., of extracting a batch-size from a tensor. Note that this extension is an alternative to introducing a separate "Dim" op (see PR: #3574 ), so only one of the two PRs is required.