Searched refs:per_shard_window_counts (Results 1 – 1 of 1) sorted by relevance
580 std::vector<int64> per_shard_window_counts(base_shape_.rank()); in ReshardAsWindowedInput() local602 per_shard_window_counts[i] = CeilOfRatio(window_count, shard_count); in ReshardAsWindowedInput()604 (wd.stride() * per_shard_window_counts[i]) % wd.base_dilation() != 0) { in ReshardAsWindowedInput()623 wd.stride() * per_shard_window_counts[i], in ReshardAsWindowedInput()626 wd.stride() * (per_shard_window_counts[i] - 1) + dilated_size; in ReshardAsWindowedInput()628 wd.stride() * per_shard_window_counts[i], in ReshardAsWindowedInput()653 wd.stride() * per_shard_window_counts[i] * shard_ordinal; in ReshardAsWindowedInput()662 wd.stride() * (per_shard_window_counts[i] - 1) + dilated_size; in ReshardAsWindowedInput()698 MultiplyAddDivideOffsetCalculation(per_shard_window_counts[i], 0, in ReshardAsWindowedInput()708 if ((wd.stride() * per_shard_window_counts[i]) % wd.base_dilation() != in ReshardAsWindowedInput()