Searched refs:dim_partitions (Results 1 – 2 of 2) sorted by relevance
297 std::vector<std::pair<int64, int64>> dim_partitions = in EmitCallToParallelForkJoin() local299 CHECK_EQ(num_partitioned_dims, dim_partitions.size()); in EmitCallToParallelForkJoin()302 const std::pair<int64, int64>& dim_partition = dim_partitions[j]; in EmitCallToParallelForkJoin()
812 int64 dim_partitions = operand_sharding.tile_assignment().dim(i); in PassthroughOperandToGatherOutputOrScatterUpdate() local815 if (dim_partitions > 1) { in PassthroughOperandToGatherOutputOrScatterUpdate()821 if (slice_size[i] != operand_shape.dimensions(i) && dim_partitions > 1) { in PassthroughOperandToGatherOutputOrScatterUpdate()830 passthrough_tile[offset_dim] = dim_partitions; in PassthroughOperandToGatherOutputOrScatterUpdate()863 int64 dim_partitions = in PassthroughGatherOutputOrScatterUpdateToOperand() local865 if (slice_size[i] != operand_shape.dimensions(i) && dim_partitions > 1) { in PassthroughGatherOutputOrScatterUpdateToOperand()873 passthrough_tile[i] = dim_partitions; in PassthroughGatherOutputOrScatterUpdateToOperand()