# ADOBE CONFIDENTIAL
# Copyright 2025 Adobe
# All Rights Reserved.
# NOTICE: All information contained herein is, and remains
# the property of Adobe and its suppliers, if any. The intellectual
# and technical concepts contained herein are proprietary to Adobe
# and its suppliers and are protected by all applicable intellectual
# property laws, including trade secret and copyright laws.
# Dissemination of this information or reproduction of this material
# is strictly forbidden unless prior written permission is obtained
# from Adobe.

from __future__ import annotations
import torch

def pad_along_last_dim(x, size):
    pad_size = size - x.shape[-1]
    if pad_size == 0:
        return x
    padding = torch.zeros(*x.shape[:-1], pad_size, dtype=x.dtype, device=x.device)
    x_padded = torch.cat([x, padding], dim=-1)
    return x_padded

def maybe_truncate_last_dim(x, size):
    if size == x.shape[-1]:
        return x
    return x[..., :size]