Add int64 shape support on GPU for stateless random ops. (#13908)
* Add int64 shape support on GPU for stateless random ops.
This fix adds int64 shape support on GPU for stateless random ops
`StatelessRandomUniform`, `StatelessRandomNormal`, `StatelessTruncatedNormal`.
The int64 shape for stateless random ops is already supported on CPU
with int32/int64 processed properly through `MakeShape`.
However, on GPU a type constraint `.TypeConstraint<int32>("T")`
has been improperly added. Such a type constraint actually prevents
an int64 shape type to run on GPU. (As a comparision, no type constraint
on CPU).
This fix removes the type constraint and allows int64 shape to be run on GPU.
This fix also adds test cases for int64 shape support on stateless random ops.
Signed-off-by:
Yong Tang <yong.tang.github@outlook.com>
* Add test cases for int64 shape support for stateless random ops.
Signed-off-by:
Yong Tang <yong.tang.github@outlook.com>
* Add int32 to shape types tested.
Loading
Please sign in to comment