Searched refs:shard_ordinal (Results 1 – 3 of 3) sorted by relevance
155 int64 Calculate(int64 shard_ordinal) const;156 HloInstruction* Calculate(HloInstruction* shard_ordinal,199 int64 Calculate(int64 shard_ordinal) const;200 HloInstruction* Calculate(HloInstruction* shard_ordinal,
678 int64 MultiplyAddDivideOffsetCalculation::Calculate(int64 shard_ordinal) const { in Calculate()679 return (shard_ordinal * multiplier_ + offset_) / divisor_; in Calculate()683 HloInstruction* shard_ordinal, SpmdBuilder* b) const { in Calculate() argument689 HloInstruction* result = shard_ordinal; in Calculate()692 scalar_shape, HloOpcode::kMultiply, shard_ordinal, in Calculate()759 int64 OffsetCalculation::Calculate(int64 shard_ordinal) const { in Calculate()762 return copy_from_.Calculate(shard_ordinal); in Calculate()764 return lhs_->Calculate(shard_ordinal) - rhs_->Calculate(shard_ordinal); in Calculate()766 return lhs_->Calculate(shard_ordinal) * rhs_->Calculate(shard_ordinal); in Calculate()772 HloInstruction* OffsetCalculation::Calculate(HloInstruction* shard_ordinal, in Calculate() argument[all …]
649 [&](int64 shard_ordinal) { in ReshardAsWindowedInput() argument650 return start_on_padded_calculations[i].Calculate(shard_ordinal) * in ReshardAsWindowedInput()653 wd.stride() * per_shard_window_counts[i] * shard_ordinal; in ReshardAsWindowedInput()659 for (int64 shard_ordinal = 0; shard_ordinal < shard_count; in ReshardAsWindowedInput() local660 ++shard_ordinal) { in ReshardAsWindowedInput()664 get_first_valid_element_offset_on_dilated_shard(shard_ordinal) + in ReshardAsWindowedInput()676 for (int64 shard_ordinal = 1; shard_ordinal < shard_count; in ReshardAsWindowedInput() local677 ++shard_ordinal) { in ReshardAsWindowedInput()679 get_first_valid_element_offset_on_dilated_shard(shard_ordinal); in ReshardAsWindowedInput()