operands_per_replica, replicated_input->getAttrs());
    replicate_op.setIsPacked(is_packed);
    operands_per_core.push_back(replicate_op);
  }

  auto pi = builder.create<TF::TPUPartitionedInputV2Op>(
      first_partitioned_input.getLoc(), replicated_input.getType(),
      operands_per_core, first_partitioned_input->getAttrs());
  pi.setIsPacked(false);  // inputs are now ops--not resources