[XLA:CPU/GPU] Fix a problem in the random number generator.

Add primitive type S64 to GetNumberOfElementsPerPhiloxRngSample so that the random number generator can generate random numbers for signed 64-bit integer. PiperOrigin-RevId: 207621810
willsim · Aug 6, 2018 · 7cc2dae · 7cc2dae
1 parent 2e7ae34
commit 7cc2dae
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tensorflow/compiler/xla/service/elemental_ir_emitter.cc b/tensorflow/compiler/xla/service/elemental_ir_emitter.cc
@@ -1302,6 +1302,7 @@ int32 GetNumberOfElementsPerPhiloxRngSample(PrimitiveType elem_prim_ty) {
     case F16:
       return 4;
     case U64:
+    case S64:
     case F64:
       return 2;
     default: