File size: 956 Bytes
300c478 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | # 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]
|