Access stateful_random_ops algorithm value as int64. Algorithm is specified to have int64 type here: https://cs.opensource.google/tensorflow/tensorflow/+/master:tensorflow/core/ops/stateful_random_ops.cc;l=38?q=stateful_random_ops.cc and below in this file.

PiperOrigin-RevId: 338326682
Change-Id: I24ba7c43258d65cd690611769e715a2616a8130f
This commit is contained in:
Anna R 2020-10-21 13:18:52 -07:00 committed by TensorFlower Gardener
parent b996deed07
commit e1b13e9e55

View File

@ -181,7 +181,7 @@ Status CompileImpl(
}
xla::Literal alg_literal;
TF_RETURN_IF_ERROR(ctx->ConstantInput(alg_input_idx, &alg_literal));
Algorithm alg = Algorithm(alg_literal.Get<int>({}));
Algorithm alg = Algorithm(alg_literal.Get<int64>({}));
if (!(alg == RNG_ALG_THREEFRY || alg == RNG_ALG_PHILOX)) {
return errors::InvalidArgument("Unsupported algorithm id: ", alg);
}