一个有条件的线性张量

Question

我有两个张量：

a = torch.nn.Parameter(torch.rand(7, requires_grad=True))
b = torch.randint(0,60, (20,))

有没有一种线性(或快速和短的方法(可以在条件下创建大小为20(类似于"b"(的张量(称为x(
即

[b<4 use a[0], 4 <=b<12 use a[1], 12<=b<22 use a[2], <28, <38, <50, >50] for every b

因此，如果：

b = [12, 93, 54, 0...]

我想要我的新张量"；x〃；将：

x = [a[2],a[6], a[6]...]

我要用这个"；x〃；张量来训练和需要值的反投影和可学习

即

loss = torch.rand(20) * x
loss.backward() ...

所以，如果其中一个a不在x中，我希望它不改变。

Answer 1

您可以对条件的乘法掩码求和：

x = a[0]*(b<4) + a[1]*((4<=b)&(b<12)) + a[2]*((12<=b)&(b<22)) + a[3]*((22<=b)&(b<28)) + a[4]*((28<=b)&(b<30)) + a[5]*((30<=b)&(b<50)) + a[6]*(b>=50)