如何用PyTorch生成满足元素上下界约束的三维随机Tensor对

首先生成你的up张量：

import torch

batch_size = 2
problem_size = 3
up = torch.randint(low=0, high=100, size=(batch_size, problem_size, problem_size))

接下来生成每个元素不超过up对应位置值的down张量，直接利用PyTorch的randint支持张量作为high参数的特性：

# 由于randint是左闭右开区间，将上限设为up+1，确保生成的最大值等于up对应元素
down = torch.randint(low=0, high=up + 1, size=(batch_size, problem_size, problem_size))

可以验证约束是否满足：

# 断言所有位置down <= up，不满足则抛出错误
assert torch.all(down <= up).item()

核心逻辑 #

• torch.randint的high参数支持传入同形状的张量，这样每个位置的随机数生成上限会对应high张量的对应值，完美实现逐位置的约束控制。
• 因为randint生成的是[low, high)区间的整数，所以必须用up + 1作为上限，这样生成的数值范围是0到up的对应元素，刚好满足down不超过up的要求。
• 就算up里有元素为0，up + 1等于1，此时randint只会生成0，完全符合约束。

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