Searched refs:PadWithValue (Results 1 – 5 of 5) sorted by relevance
274 lhs = lhs.Reshard(aligned_lhs_sharding).PadWithValue(zero); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()275 rhs = rhs.PadWithValue(zero); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()280 lhs = lhs.PadWithValue(zero); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()281 rhs = rhs.Reshard(aligned_rhs_sharding).PadWithValue(zero); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()583 lhs = lhs.Reshard(aligned_lhs_sharding).PadWithValue(zero); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()584 rhs = rhs.PadWithValue(zero, reversed_rhs_dims); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()589 lhs = lhs.PadWithValue(zero); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()591 rhs.Reshard(aligned_rhs_sharding).PadWithValue(zero, reversed_rhs_dims); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()
622 to_mask.PadWithValue(b->AddInstruction(HloInstruction::CreateConstant( in PartitionBaseCase()628 lhs = lhs.PadWithValue(zero); in PartitionBaseCase()629 rhs = rhs.PadWithValue(zero); in PartitionBaseCase()1504 lhs.Reshard(*rhs_sharding_transposed_to_match_lhs).PadWithValue(zero); in PartitionBaseCase()1505 rhs = rhs.PadWithValue(zero); in PartitionBaseCase()1507 lhs = lhs.PadWithValue(zero); in PartitionBaseCase()1509 rhs.Reshard(*lhs_sharding_transposed_to_match_rhs).PadWithValue(zero); in PartitionBaseCase()1613 lhs.Reshard(*rhs_sharding_transposed_to_match_lhs).PadWithValue(zero); in PartitionBaseCase()1614 rhs = rhs.PadWithValue(zero); in PartitionBaseCase()1616 lhs = lhs.PadWithValue(zero); in PartitionBaseCase()[all …]
372 .PadWithValue(CreateR0WithType(hlo->shape().element_type(), 0, &b_)); in HandleFft()
313 PartitionedHlo PadWithValue(HloInstruction* pad_value,
479 PartitionedHlo PartitionedHlo::PadWithValue( in PadWithValue() function in xla::spmd::PartitionedHlo1671 auto partitioned_input = GetPartitionedHlo(input).PadWithValue( in HandleSort()1677 .PadWithValue(CreateLastWithType(index_type, &b_)); in HandleSort()1780 input_partitioned = input_partitioned.PadWithValue( in HandleCustomCall()1834 auto partitioned_input = GetPartitionedHlo(input).PadWithValue( in HandleCustomCall()2730 inputs.back().PadWithValue(inits[operand_id], /*left_padded_dims=*/{}, in HandleReduce()