Jax中padding的逆操作实现：如何支持类似PyTorch的负padding功能

PyTorch的F.pad支持负数padding的核心逻辑是：正padding值为边缘填充对应长度的元素，负padding值为边缘裁剪对应长度的元素。Jax原生的jnp.pad仅支持非负填充值，我们只需要额外封装一层处理正负参数的逻辑即可对齐PyTorch的用法。

封装代码 #

import jax.numpy as jnp

def pad_like_torch(x, padding, mode="constant", constant_values=0):
    # 适配PyTorch的padding参数顺序：从最后一维开始给出前后填充值
    padding = jnp.asarray(padding, dtype=int).reshape(-1, 2)[::-1]
    # 补全所有维度的padding参数，未指定维度默认填充0
    full_pad = jnp.concatenate([
        jnp.zeros((x.ndim - len(padding), 2), dtype=int),
        padding
    ], axis=0)

    slices = []
    valid_pad = []
    for dim_idx, (pad_before, pad_after) in enumerate(full_pad):
        # 处理负padding：对应维度裁剪
        slice_start = max(-pad_before, 0)
        slice_end = x.shape[dim_idx] - max(-pad_after, 0)
        slices.append(slice(slice_start, slice_end if slice_end > slice_start else 0))
        # 提取正padding参数传给原生jnp.pad
        valid_pad.append((max(pad_before, 0), max(pad_after, 0)))
    
    # 先执行裁剪逻辑
    x_cropped = x[tuple(slices)]
    # 再执行正填充逻辑
    if mode == "constant":
        return jnp.pad(x_cropped, valid_pad, mode=mode, constant_values=constant_values)
    return jnp.pad(x_cropped, valid_pad, mode=mode)

效果验证 #

你提到的F.pad(array, [-1,-1])调用效果，对应实现如下：

# 测试输入
arr = jnp.array([[1,2,3,4], [5,6,7,8]])

# 最后一维前后各裁剪1位，完全对齐PyTorch用法
output = pad_like_torch(arr, [-1, -1])
print(output)
# 输出结果：
# [[2 3]
#  [6 7]]

注意事项 #

• 支持所有jnp.pad原生兼容的填充模式，包括constant、edge、reflect等，和PyTorch对应模式输出一致
• 多维度padding场景完全对齐PyTorch传参规则，例如对倒数第二维裁剪顶部2位、底部1位，最后一维左侧填充2位：pad_like_torch(arr, [2, 0, -2, -1])

内容的提问来源于stack exchange，提问作者JupiterNo41