| # 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] | |