Searched refs:new_all_to_all_shape (Results 1 – 1 of 1) sorted by relevance
66 Shape new_all_to_all_shape; in ExpandInstruction() local67 new_all_to_all_shape.set_element_type( in ExpandInstruction()71 new_all_to_all_shape.add_dimensions(all_to_all->shape().dimensions(i)); in ExpandInstruction()74 new_all_to_all_shape.add_dimensions(all_to_all_group_size); in ExpandInstruction()75 new_all_to_all_shape.add_dimensions(split_size); in ExpandInstruction()77 new_all_to_all_shape.add_dimensions(1); in ExpandInstruction()80 *(new_all_to_all_shape.mutable_layout()) = in ExpandInstruction()84 new_all_to_all_shape, instruction->mutable_operand(0))); in ExpandInstruction()88 new_all_to_all_shape, {operand_reshape})); in ExpandInstruction()