Add int64 support of `axis` (`Tidx`) for ConcatV2
In `array_ops.cc`, it was specified that ConcatV2 support
both int32 and int64 data types of `axis` (`Tidx`):
```
.Attr("Tidx: {int32, int64} = DT_INT32")
```
However, in actual kernel implementations only int32 is supported
as there is an unnecessary `.TypeConstraint<int32>("Tidx")` specified.
This fix tries to address the discrepancy between the ops declaration
and kernel registration by adding the int64 axis (`Tidx`) support
for `ConcatV2`.
This fix removes the TypeConstraint and adds additional processing
so that differnt types (int32 or int64) of `axis` could be processed
correctly.
Additional test cases have been added to cover the changes as well.
Signed-off-by:
Yong Tang <yong.tang.github@outlook.com>
Loading
Please sign in to comment