提交 9fdaff45 编写于 作者: M Mihai Maruseac

Prevent `int64` to `int` truncation in `Shard` API usage.

The function argument in `Shard` must be a function of two `int64` arguments. However, we are passing in a function with two `int` arguments. Thus, for large workloads, these arguments get truncated from positive `int64` values to negative `int` ones, resulting in a buffer out of bounds write.

PiperOrigin-RevId: 332557334
Change-Id: I236c9a2e7f53580e520571da8ba941a3aa9fa0b5
上级 01948c69
......@@ -204,7 +204,7 @@ class RandomGammaOp : public OpKernel {
// avoid a couple flops which can be done on a per-alpha basis.
auto DoWork = [num_samples, num_alphas, &rng, samples_flat, alpha_flat](
int start_output, int limit_output) {
int64 start_output, int64 limit_output) {
using Eigen::numext::exp;
using Eigen::numext::log;
using Eigen::numext::pow;
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册