Layout flow

class LayoutFlowConfig(ConfigMixin):
    """Configuration saved with converted LayoutFlow pipelines."""

    config_name: str = "layout_flow_config.json"

    @register_to_config
    def __init__(
        self,
        *,
        dataset_name: str = "publaynet",
        id2label: dict[int | str, str] | None = None,
        max_length: int = 20,
        latent_dim: int = 128,
        d_model: int = 512,
        nhead: int = 8,
        dim_feedforward: int = 2048,
        num_layers: int = 4,
        dropout: float = 0.1,
        use_pos_enc: bool = False,
        tr_enc_only: bool = True,
        attr_encoding: AttrEncoding = AttrEncoding.analog_bit,
        seq_type: SeqType = SeqType.stacked,
        distribution: InitialDistributionName = InitialDistributionName.gaussian,
        sample_padding: bool = False,
        inference_steps: int = 100,
        ode_solver: OdeSolverName = OdeSolverName.euler,
        bbox_format: BoxFormat | str = "xywh",
        coordinate_range: CoordinateRange = CoordinateRange.normalized_0_1,
    ) -> None:
        """Initialize LayoutFlow pipeline and model settings.

        Args:
            dataset_name: Dataset variant or alias.
            id2label: Optional explicit id-to-label mapping.
            max_length: Maximum number of layout elements.
            latent_dim: Latent dimension.
            d_model: Transformer hidden size.
            nhead: Number of attention heads.
            dim_feedforward: Feed-forward hidden size.
            num_layers: Number of transformer layers.
            dropout: Dropout probability.
            use_pos_enc: Whether to add sinusoidal position encodings.
            tr_enc_only: Whether to use the encoder-only path.
            attr_encoding: Attribute encoding used by the checkpoint.
            seq_type: Sequence layout type.
            distribution: Initial sampling distribution.
            sample_padding: Whether sampling includes padded elements.
            inference_steps: Default Euler inference steps.
            ode_solver: ODE solver name.
            bbox_format: Public bounding-box format.
            coordinate_range: Public coordinate range.

        Raises:
            ValueError: If ``dataset_name`` is unsupported.
        """
        self.dataset_name = str(normalize_dataset_name(dataset_name))
        raw_id2label = id2label or default_id2label(self.dataset_name)
        self.id2label = {int(k): v for k, v in raw_id2label.items()}

        self.max_length = max_length
        self.latent_dim = latent_dim
        self.d_model = d_model
        self.nhead = nhead
        self.dim_feedforward = dim_feedforward
        self.num_layers = num_layers
        self.dropout = dropout
        self.use_pos_enc = use_pos_enc
        self.tr_enc_only = tr_enc_only

        self.attr_encoding = str(AttrEncoding(attr_encoding))
        self.seq_type = str(SeqType(seq_type))
        self.distribution = str(InitialDistributionName(distribution))
        self.sample_padding = sample_padding

        self.inference_steps = inference_steps
        self.ode_solver = str(OdeSolverName(ode_solver))
        self.bbox_format = str(BoxFormat(bbox_format))
        self.coordinate_range = str(CoordinateRange(coordinate_range))

    @property
    def label2id(self) -> dict[str, int]:
        """Return the label-name to integer-id mapping."""
        return {v: k for k, v in self.id2label.items()}

    @property
    def num_labels(self) -> int:
        """Return the number of labels, including the background label."""
        return len(self.id2label)

    @property
    def attr_dim(self) -> int:
        """Return the analog-bit attribute dimensionality."""
        if AttrEncoding(self.attr_encoding) is AttrEncoding.analog_bit:
            return int(math.ceil(math.log2(self.num_labels)))
        return 1

    @property
    def sample_dim(self) -> int:
        """Return the model-state dimensionality per layout element."""
        return 4 + self.attr_dim

@dataclass
class LayoutFlowModelOutput(BaseOutput):
    """Output of ``LayoutFlowTransformerModel``."""

    sample: Float[torch.Tensor, "batch elements channels"]

class LayoutFlowTransformerModel(ModelMixin, ConfigMixin):
    """Diffusers model wrapper around the LayoutFlow backbone."""

    config_name: str = "layout_flow_model_config.json"

    @register_to_config
    def __init__(
        self,
        *,
        num_labels: int = 6,
        latent_dim: int = 128,
        tr_enc_only: bool = True,
        d_model: int = 512,
        nhead: int = 8,
        dim_feedforward: int = 2048,
        num_layers: int = 4,
        dropout: float = 0.1,
        use_pos_enc: bool = False,
        attr_encoding: AttrEncoding = AttrEncoding.analog_bit,
        seq_type: SeqType = SeqType.stacked,
    ) -> None:
        """Initialize the converted LayoutFlow transformer model.

        Args:
            num_labels: Number of dataset labels.
            latent_dim: Latent dimension.
            tr_enc_only: Whether to use the encoder-only path.
            d_model: Transformer hidden size.
            nhead: Number of attention heads.
            dim_feedforward: Feed-forward hidden size.
            num_layers: Number of transformer layers.
            dropout: Dropout probability.
            use_pos_enc: Whether to add positional encodings.
            attr_encoding: Attribute encoding.
            seq_type: Sequence type.
        """
        super().__init__()
        self.geom_dim = 4
        self.attr_dim = (
            int(np.ceil(np.log2(num_labels)))
            if AttrEncoding(attr_encoding) is AttrEncoding.analog_bit
            else 1
        )
        self.backbone = LayoutDMBackbone(
            latent_dim=latent_dim,
            tr_enc_only=tr_enc_only,
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=dim_feedforward,
            num_layers=num_layers,
            dropout=dropout,
            use_pos_enc=use_pos_enc,
            num_cat=num_labels,
            attr_encoding=attr_encoding,
            seq_type=seq_type,
        )

    def forward(
        self,
        sample: Float[torch.Tensor, "batch elements channels"],
        timestep: Float[torch.Tensor, ""],
        cond_mask: Bool[torch.Tensor, "batch elements channels"],
        return_dict: bool = True,
    ) -> LayoutFlowModelOutput | tuple[Float[torch.Tensor, "batch elements channels"]]:
        """Predict the vector field for a model state.

        Args:
            sample: Current model state.
            timestep: Current integration timestep.
            cond_mask: Condition mask.
            return_dict: Whether to return a dataclass output.

        Returns:
            Model output dataclass or single-item tuple.
        """
        timestep = timestep.to(device=sample.device, dtype=sample.dtype)
        if timestep.ndim == 0:
            timestep = timestep.repeat(sample.shape[0])
        geom = sample[:, :, : self.geom_dim]
        attr = sample[:, :, self.geom_dim :]
        out = self.backbone(geom, attr, cond_mask.to(torch.long), timestep)
        if not return_dict:
            return (out,)
        return LayoutFlowModelOutput(sample=out)

class LayoutFlowPipeline(DiffusionPipeline):
    """Generate layouts with a converted LayoutFlow checkpoint."""

    model_cpu_offload_seq: str = "model"

    def __init__(
        self,
        model: LayoutFlowTransformerModel,
        scheduler: LayoutFlowEulerScheduler,
        config: LayoutFlowConfig,
        processor: LayoutFlowProcessor | None = None,
    ) -> None:
        """Create a LayoutFlow pipeline.

        Args:
            model: Converted LayoutFlow transformer model.
            scheduler: Increasing-time Euler scheduler.
            config: Pipeline configuration.
            processor: Optional input/output processor.
        """
        super().__init__()
        self.register_modules(model=model, scheduler=scheduler)
        self.layout_flow_config = config
        self.processor = processor or LayoutFlowProcessor(self.layout_flow_config)
        self.model.eval()

    @torch.no_grad()
    def __call__(
        self,
        *,
        batch_size: int = 1,
        seed: int | None = None,
        generator: torch.Generator | None = None,
        condition_type: ConditionType | str = ConditionType.unconditional,
        labels: Int[torch.Tensor, "batch elements"]
        | Int[np.ndarray, "batch elements"]
        | Sequence[ArrayLikeInput]
        | None = None,
        bbox: Float[torch.Tensor, "batch elements 4"]
        | Float[np.ndarray, "batch elements 4"]
        | Sequence[ArrayLikeInput]
        | None = None,
        mask: Bool[torch.Tensor, "batch elements"]
        | Bool[np.ndarray, "batch elements"]
        | Sequence[ArrayLikeInput]
        | None = None,
        num_elements: int | list[int] | Int[torch.Tensor, "batch"] | None = None,
        box_format: BoxFormat | str = "xywh",
        normalized: bool = True,
        canvas_size: tuple[int, int] | None = None,
        num_inference_steps: int | None = None,
        guidance_scale: float = 0.0,
        output_type: OutputType | str = "dataclass",
        return_intermediates: bool = False,
    ) -> LayoutGenerationOutput | dict[str, Shaped[torch.Tensor, "..."]]:
        """Generate layout boxes and labels.

        Args:
            batch_size: Number of layouts to generate.
            seed: Optional seed used when ``generator`` is omitted.
            generator: Optional torch random generator.
            condition_type: Public condition name or supported alias.
            labels: Optional condition labels.
            bbox: Optional condition boxes.
            mask: Optional valid-element mask.
            num_elements: Optional element counts for unconditional masks.
            box_format: Input and output box format.
            normalized: Whether coordinates are normalized.
            canvas_size: Pixel canvas size for denormalized coordinates.
            num_inference_steps: Number of Euler steps.
            guidance_scale: Classifier-free guidance scale.
            output_type: ``"dataclass"`` or ``"dict"``.
            return_intermediates: Whether to include intermediate samples.

        Returns:
            Layout generation output dataclass, or a dictionary when requested.

        Raises:
            ValueError: If ``condition_type``, ``box_format``, or ``output_type``
                is unsupported.

        Examples:
            >>> pipe = LayoutFlowPipeline(
            ...     model=LayoutFlowTransformerModel(
            ...         num_labels=6, latent_dim=8, d_model=16, nhead=4,
            ...         dim_feedforward=32, num_layers=1
            ...     ),
            ...     scheduler=LayoutFlowEulerScheduler(num_inference_steps=2),
            ...     config=LayoutFlowConfig(max_length=2, latent_dim=8, d_model=16),
            ... )
            >>> out = pipe(batch_size=1, num_elements=1, seed=0, num_inference_steps=2)
            >>> out.bbox.shape[-1]
            4
        """
        if generator is None and seed is not None:
            generator = torch.Generator(device=self.device).manual_seed(seed)
        canonical = normalize_condition_type(condition_type)
        output_kind = OutputType(output_type)
        processed = self.processor(
            bbox=bbox,
            labels=labels,
            mask=mask,
            num_elements=num_elements,
            batch_size=batch_size,
            box_format=box_format,
            normalized=normalized,
            canvas_size=canvas_size,
            device=self.device,
        )
        batch_size = processed["bbox"].shape[0]
        cond_mask = self.processor.make_condition_mask(
            canonical, mask=processed["mask"], generator=generator
        )
        cond_state = self.processor.preprocess_state(
            self.processor.model_state(processed["bbox"], processed["labels"])
        )
        sample = sample_initial_state(
            batch_size=batch_size,
            max_length=self.layout_flow_config.max_length,
            lengths=processed["length"],
            dim=self.layout_flow_config.sample_dim,
            distribution=self.layout_flow_config.distribution,
            generator=generator,
            device=self.device,
            dtype=cond_state.dtype,
        )
        if canonical is ConditionType.refinement:
            sample = cond_state
            self.scheduler.set_timesteps(
                num_inference_steps, device=self.device, start=0.97, end=1.0
            )
        else:
            self.scheduler.set_timesteps(num_inference_steps, device=self.device)
        trajectory = [] if return_intermediates else None
        for i, timestep in enumerate(self.scheduler.timesteps[:-1]):
            x_in = (1 - cond_mask) * cond_state + cond_mask * sample
            t_batch = timestep.repeat(batch_size)
            vector = self.model(
                sample=x_in, timestep=t_batch, cond_mask=cond_mask
            ).sample
            if guidance_scale:
                uncond_mask = torch.ones_like(cond_mask)
                uncond = self.model(
                    sample=x_in, timestep=t_batch, cond_mask=uncond_mask
                ).sample
                vector = (1 + guidance_scale) * vector - guidance_scale * uncond
            sample = self.scheduler.step(
                vector,
                timestep,
                sample,
                next_timestep=self.scheduler.timesteps[i + 1],
            ).prev_sample
            if trajectory is not None:
                trajectory.append(sample.detach().cpu())
        final_state = (1 - cond_mask) * cond_state + cond_mask * sample
        decoded = self.processor.postprocess(
            final_state,
            mask=processed["mask"],
            box_format=box_format,
            normalized=normalized,
            canvas_size=canvas_size,
        )
        output = LayoutGenerationOutput(
            bbox=decoded["bbox"].detach().cpu(),
            labels=decoded["labels"].detach().cpu(),
            mask=decoded["mask"].detach().cpu(),
            id2label=self.layout_flow_config.id2label,
            trajectory=trajectory,
            intermediates={"condition_type": str(canonical)}
            if return_intermediates
            else None,
        )
        if output_kind is OutputType.dict:
            return dict(output)
        if output_kind is OutputType.dataclass:
            return output
        assert_never(output_kind)

    generate = __call__

    def save_pretrained(self, save_directory: str | Path) -> None:
        """Save pipeline components and LayoutFlow config.

        Args:
            save_directory: Output directory.
        """
        super().save_pretrained(save_directory)
        self.layout_flow_config.save_config(save_directory)

    @classmethod
    def from_pretrained(
        cls, pretrained_model_name_or_path: str | Path
    ) -> LayoutFlowPipeline:
        """Load a saved LayoutFlow pipeline.

        Args:
            pretrained_model_name_or_path: Local directory or Hub id.

        Returns:
            Loaded LayoutFlow pipeline.
        """
        config_dict, _ = LayoutFlowConfig.load_config(
            pretrained_model_name_or_path,
            return_unused_kwargs=True,
        )
        config = cast(LayoutFlowConfig, LayoutFlowConfig.from_config(config_dict))
        pipe = super().from_pretrained(pretrained_model_name_or_path, config=config)
        pipe.layout_flow_config = config
        pipe.processor = LayoutFlowProcessor(config)
        return pipe

class ConditionType(StrEnum):
    """Canonical condition names used by layout generation interfaces."""

    unconditional = auto()
    label = auto()
    label_size = auto()
    completion = auto()
    refinement = auto()
    text = auto()
    content_image = auto()
    relation = auto()
    hierarchical = auto()
    retrieval = auto()

class LayoutFlowProcessor(ProcessorMixin):
    """Prepare public layout tensors for the LayoutFlow model."""

    config_name = "processor_config.json"

    def __init__(self, config: LayoutFlowConfig) -> None:
        """Create a processor for a LayoutFlow configuration.

        Args:
            config: LayoutFlow pipeline configuration.
        """
        super().__init__()
        self.config = config
        bit_mask = [1 << k for k in range(config.attr_dim)]
        self.bit_mask = torch.tensor(bit_mask, dtype=torch.long)

    def __call__(
        self,
        *,
        bbox: Float[torch.Tensor, "batch elements 4"]
        | Float[np.ndarray, "batch elements 4"]
        | Sequence[ArrayLikeInput]
        | None = None,
        labels: Int[torch.Tensor, "batch elements"]
        | Int[np.ndarray, "batch elements"]
        | Sequence[ArrayLikeInput]
        | None = None,
        mask: Bool[torch.Tensor, "batch elements"]
        | Bool[np.ndarray, "batch elements"]
        | Sequence[ArrayLikeInput]
        | None = None,
        num_elements: int | list[int] | Int[torch.Tensor, "batch"] | None = None,
        batch_size: int = 1,
        box_format: BoxFormat | str = "xywh",
        normalized: bool = True,
        canvas_size: tuple[int, int] | None = None,
        device: torch.device | str | None = None,
    ) -> dict[str, Shaped[torch.Tensor, "..."]]:
        """Convert public inputs into padded model-ready tensors.

        Args:
            bbox: Optional boxes in the requested ``box_format``.
            labels: Optional dataset-local label ids.
            mask: Optional valid-element mask.
            num_elements: Optional element counts used when ``mask`` is omitted.
            batch_size: Batch size used when tensors are omitted.
            box_format: Format of ``bbox``.
            normalized: Whether ``bbox`` is already normalized to ``[0, 1]``.
            canvas_size: Pixel canvas size required for denormalized boxes.
            device: Target torch device.

        Returns:
            Dictionary with ``bbox``, ``labels``, ``mask``, and ``length`` tensors.

        Raises:
            ValueError: If denormalized boxes are missing ``canvas_size`` or the
                box format is unsupported.
        """
        device = torch.device(device) if device is not None else torch.device("cpu")
        max_length = self.config.max_length
        if bbox is None:
            bbox_t = torch.zeros(
                batch_size, max_length, 4, dtype=torch.float32, device=device
            )
        else:
            bbox_t = torch.as_tensor(bbox, dtype=torch.float32, device=device)
            if bbox_t.ndim == 2:
                bbox_t = bbox_t.unsqueeze(0)
            batch_size = bbox_t.shape[0]
            if not normalized:
                if canvas_size is None:
                    raise ValueError("canvas_size is required when normalized=False")

                bbox_t = normalize_boxes(
                    bbox_t, canvas_size=canvas_size, box_format=box_format
                )
            else:
                fmt = BoxFormat(box_format)
                if fmt is BoxFormat.ltwh:
                    bbox_t = ltwh_to_xywh(bbox_t)
                elif fmt is BoxFormat.ltrb:
                    bbox_t = ltrb_to_xywh(bbox_t)
                elif fmt is not BoxFormat.xywh:
                    assert_never(fmt)
            bbox_t = self._pad_tensor(bbox_t, max_length, 0.0)
        if labels is None:
            labels_t = torch.zeros(
                batch_size, max_length, dtype=torch.long, device=device
            )
        else:
            labels_t = torch.as_tensor(labels, dtype=torch.long, device=device)
            if labels_t.ndim == 1:
                labels_t = labels_t.unsqueeze(0)
            labels_t = self._pad_tensor(labels_t, max_length, 0)
        if mask is None:
            lengths = self._num_elements_to_lengths(
                num_elements, batch_size, max_length, device
            )
            mask_t = torch.arange(max_length, device=device)[None, :] < lengths[:, None]
        else:
            mask_t = torch.as_tensor(mask, dtype=torch.bool, device=device)
            if mask_t.ndim == 1:
                mask_t = mask_t.unsqueeze(0)
            mask_t = self._pad_tensor(mask_t, max_length, False)
            lengths = mask_t.sum(dim=1).long()
        bbox_t = bbox_t * mask_t.unsqueeze(-1)
        labels_t = labels_t * mask_t.long()
        return {"bbox": bbox_t, "labels": labels_t, "mask": mask_t, "length": lengths}

    def encode_labels(
        self, labels: Int[torch.Tensor, "batch elements"]
    ) -> Float[torch.Tensor, "batch elements bits"]:
        """Encode integer labels as analog-bit vectors."""
        bit_mask = self.bit_mask.to(labels.device)
        return (
            torch.bitwise_and(labels.unsqueeze(-1), bit_mask).float() / bit_mask.float()
        )

    def decode_labels(
        self, bits: Float[torch.Tensor, "batch elements bits"]
    ) -> Int[torch.Tensor, "batch elements"]:
        """Decode analog-bit vectors into integer labels."""
        bit_mask = self.bit_mask.to(bits.device)
        active = (bits - 0.5 >= 0).long()
        return (
            (active * bit_mask).sum(dim=-1).clamp(0, self.config.num_labels - 1).long()
        )

    def model_state(
        self,
        bbox: Float[torch.Tensor, "batch elements 4"],
        labels: Int[torch.Tensor, "batch elements"],
    ) -> Float[torch.Tensor, "batch elements channels"]:
        """Concatenate normalized boxes and analog-bit labels."""
        return torch.cat([bbox, self.encode_labels(labels)], dim=-1)

    def preprocess_state(
        self,
        state: Float[torch.Tensor, "batch elements channels"],
        *,
        reverse: bool = False,
    ) -> Float[torch.Tensor, "batch elements channels"]:
        """Map between public ``[0, 1]`` state and model distribution range."""
        if self.config.distribution in {"gaussian", "gmm", "uniform", "gauss_uniform"}:
            return (state + 1) / 2 if reverse else 2 * state - 1
        return state

    def make_condition_mask(
        self,
        condition_type: ConditionType,
        *,
        mask: Bool[torch.Tensor, "batch elements"],
        generator: torch.Generator | None = None,
    ) -> Int[torch.Tensor, "batch elements channels"]:
        """Create the condition mask for a conditioning mode.

        Args:
            condition_type: Canonical condition or alias.
            mask: Valid-element mask.
            generator: Optional generator used by completion masking.

        Returns:
            Long tensor where ``1`` means generated and ``0`` means conditioned.

        Raises:
            ValueError: If the condition type is unsupported.
        """
        batch, seq = mask.shape
        cond_mask = torch.ones(
            batch,
            seq,
            self.config.sample_dim,
            dtype=torch.long,
            device=mask.device,
        )
        if condition_type in {ConditionType.label, ConditionType.refinement}:
            cond_mask[:, :, 4:] = 0
        elif condition_type is ConditionType.label_size:
            cond_mask[:, :, 2:] = 0
        elif condition_type is ConditionType.completion:
            cond_mask = self._completion_mask(cond_mask, mask, generator)
        elif condition_type is ConditionType.unconditional:
            pass
        else:
            raise ValueError(f"Unsupported LayoutFlow condition_type: {condition_type}")

        return cond_mask

    def postprocess(
        self,
        state: Float[torch.Tensor, "batch elements channels"],
        *,
        mask: Bool[torch.Tensor, "batch elements"],
        box_format: BoxFormat | str = "xywh",
        normalized: bool = True,
        canvas_size: tuple[int, int] | None = None,
    ) -> dict[str, Shaped[torch.Tensor, "..."]]:
        """Convert model state back to public layout tensors.

        Args:
            state: Model state tensor.
            mask: Valid-element mask.
            box_format: Requested output box format.
            normalized: Whether to return normalized coordinates.
            canvas_size: Pixel canvas size for denormalized coordinates.

        Returns:
            Dictionary with ``bbox``, ``labels``, and ``mask``.

        Raises:
            ValueError: If denormalized output is requested without
                ``canvas_size``.
        """
        restored = self.preprocess_state(state, reverse=True)
        bbox = clamp_boxes(restored[:, :, :4]) * mask.unsqueeze(-1)
        labels = self.decode_labels(restored[:, :, 4:]) * mask.long()
        fmt = BoxFormat(box_format)
        if fmt is not BoxFormat.xywh:
            if not normalized and canvas_size is None:
                raise ValueError("canvas_size is required for denormalized output")

            if canvas_size is None:
                canvas_size = (1, 1)
            bbox = denormalize_boxes(bbox, canvas_size=canvas_size, box_format=fmt)
            if normalized:
                scale = torch.tensor(
                    (*canvas_size, *canvas_size), dtype=bbox.dtype, device=bbox.device
                )
                bbox = bbox / scale
        elif not normalized:
            if canvas_size is None:
                raise ValueError("canvas_size is required for denormalized output")

            bbox = denormalize_boxes(
                bbox, canvas_size=canvas_size, box_format=BoxFormat.xywh
            )
        return {"bbox": bbox, "labels": labels, "mask": mask}

    def _completion_mask(
        self,
        cond_mask: Int[torch.Tensor, "batch elements channels"],
        mask: Bool[torch.Tensor, "batch elements"],
        generator: torch.Generator | None,
    ) -> Int[torch.Tensor, "batch elements channels"]:
        for i, length in enumerate(mask.sum(dim=1).tolist()):
            if length <= 1:
                continue
            keep = max(1, int(length * 0.2))
            scores = torch.rand(length, device=mask.device, generator=generator)
            idx = scores.topk(keep).indices
            cond_mask[i, idx] = 0
        return cond_mask

    @staticmethod
    def _pad_tensor(
        tensor: Shaped[torch.Tensor, "..."], max_length: int, value: float | int | bool
    ) -> Shaped[torch.Tensor, "..."]:
        if tensor.shape[1] > max_length:
            return tensor[:, :max_length]
        if tensor.shape[1] == max_length:
            return tensor
        pad_shape = (tensor.shape[0], max_length - tensor.shape[1], *tensor.shape[2:])
        pad = torch.full(pad_shape, value, dtype=tensor.dtype, device=tensor.device)
        return torch.cat([tensor, pad], dim=1)

    @staticmethod
    def _num_elements_to_lengths(
        num_elements: int | list[int] | Int[torch.Tensor, "batch"] | None,
        batch_size: int,
        max_length: int,
        device: torch.device,
    ) -> Int[torch.Tensor, "batch"]:
        if num_elements is None:
            return torch.full(
                (batch_size,), max_length, dtype=torch.long, device=device
            )
        lengths = torch.as_tensor(num_elements, dtype=torch.long, device=device)
        if lengths.ndim == 0:
            lengths = lengths.repeat(batch_size)
        return lengths.clamp(0, max_length)