Hi!
I don’t know any compact way to achieve this, however, using the function torch.mul might be useful for this task. Here’s a possible solution:
A = torch.randn(5,5)
B = torch.randint(1,4,(5,5)).float()
Z = torch.zeros(5,5).float()

# I_n is a matrix where I_n(i,j) = 1 iff B(i,j)=n and for else I_n(i,j)=0
I1 = torch.where(B==1, B, Z)
I2 = torch.where(B==2, torch.ones(5,5), Z)
I3 = torch.where(B==3, torch.ones(5,5), Z)

# observe that torch.mul is different than torch.mm
step_1 = torch.sum(torch.mul(I1, A),dim=1)

# reshape for transposing
pre_step_2 = step_1.reshape(-1,1).repeat(1,5)
step_2 = torch.mul(torch.mul(I2, A), pre_step_2)

# here we also sum the components which are not corresponding to 2,
# which are 0, so they will become 1 under exp 
step_3 = torch.sum(torch.exp(step_2))

result = step_3-torch.sum(I1+I3)