Diff Coverage

        self.sequence_parallel = bool(mesh_context.sequence_parallel)
        self.runtime_input_adapter = runtime_input_adapter
        self.cp_sharder = CPBatchSharder(self.parallel_context, token_fields=OMNI_CP_TOKEN_FIELDS)
        self.tp_broadcaster = TPBatchBroadcaster(self.parallel_context, device)
        self.batch_transport = None
        self._batch_index = 0

    def __call__(
            self,

                tuple(local_batch["input_ids"].shape), self.parallel_context.tp_world_size,
                self.parallel_context.cp_world_size, self.parallel_context.cp_rank, enabled=True,
            )

        transport_inputs = {}
        if self.batch_transport is None:
            device_batch = self.tp_broadcaster.broadcast(local_batch, broadcast_all_fields=True)
        else:
            device_batch, transport_inputs = self.batch_transport(local_batch, self.tp_broadcaster)
        runtime_inputs = self._build_runtime_inputs(device_batch)
        model_inputs, loss_inputs = self._split_model_and_loss_inputs(device_batch)
        if set(transport_inputs).intersection((*device_batch, *runtime_inputs)):
            raise ValueError("batch transport cannot replace Omni model or runtime inputs")
        model_inputs.update(runtime_inputs)
        model_inputs.update(transport_inputs)
        # Shapes are host metadata; inspecting device tensor values here can stall collectives.
        logger.debug(
            "Omni get_batch: batch=%d transport_ready fields=%s runtime_fields=%s",
            self._batch_index,

    def __init__(self, *, runtime: EncoderParallelRuntime, chunk_bytes: int = _CHUNK_BYTES) -> None:
        """Set the runtime and maximum device staging size in bytes."""
        if isinstance(chunk_bytes, bool) or not isinstance(chunk_bytes, int) or chunk_bytes < 1:
            raise ValueError("encoder transport chunk_bytes must be a positive integer")
        self.runtime = runtime
        self.chunk_bytes = chunk_bytes

    def _prepare_source(self, batch: Mapping[str, Any]) -> dict[str, Any]:

                raise ValueError("encoder transport fields cannot be supplied by the DataLoader")
            for name in bridge.input_fields:
                value = batch.get(name)
                if torch.is_tensor(value) and value.device.type != "cpu":
                    raise ValueError(f"encoder source field {name} must remain on the CPU")
            items = bridge.parse_images(batch)
            pixels = batch.get(bridge.payload_field)
            if pixels is not None:
                dtype = pixels.dtype

    def _check_source_errors(self, error: str | None, device: Any) -> None:
        """Make all FSDP participants fail before plan or payload collectives."""
        invalid = torch.tensor(int(error is not None), device=device)
        if dist.is_initialized() and self.runtime.global_mesh.size() > 1:
            dist.all_reduce(invalid, op=dist.ReduceOp.MAX, group=self.runtime.global_mesh.get_group())
        if bool(invalid.item()):
            raise ValueError(f"invalid encoder CPU input: {error or 'another source rank rejected its batch'}")

    def __call__(self, batch: Mapping[str, Any] | None,

        )
        if encoder_runtime is not None:
            if not isinstance(get_batch, OmniParallelBatch):
                raise ValueError("encoder parallel requires OmniParallelBatch")
            get_batch.batch_transport = EncoderParallelBatchTransport(runtime=encoder_runtime)
        self.base.get_batch = get_batch

    @property
    def distributed_setup(self) -> Any: