Commit ac0004e7 authored by Yong Tang's avatar Yong Tang Committed by Vijay Vasudevan
Browse files

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: default avatarYong Tang <yong.tang.github@outlook.com>

* Add test cases for int64 shape support for stateless random ops.

Signed-off-by: default avatarYong Tang <yong.tang.github@outlook.com>

* Add int32 to shape types tested.
parent 0d437c3b
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment