Searched refs:input_shard_size (Results 1 – 2 of 2) sorted by relevance
804 int64 input_shard_size = hlo->shape().dimensions(dim); in ExchangeHalo() local812 if (max_left_halo_size + max_right_halo_size + input_shard_size >= in ExchangeHalo()813 input_shard_size * shard_count && in ExchangeHalo()814 (max_left_halo_size > input_shard_size || in ExchangeHalo()815 max_right_halo_size > input_shard_size)) { in ExchangeHalo()819 for (int64 i = CeilOfRatio(max_left_halo_size, input_shard_size) - 1; i >= 0; in ExchangeHalo()832 std::min(max_left_halo_size - input_shard_size * i, input_shard_size); in ExchangeHalo()853 for (int64 i = 0; i < CeilOfRatio(max_right_halo_size, input_shard_size); in ExchangeHalo()866 std::min(max_right_halo_size - input_shard_size * i, input_shard_size); in ExchangeHalo()
787 int64 input_shard_size = in ReshardAsWindowedInput() local794 input_shard_size, explicit_left_padding[dim], 1); in ReshardAsWindowedInput()800 input_shard_size, input_shard_size + explicit_left_padding[dim], 1); in ReshardAsWindowedInput()