Torch Operations Using Custom Functions

I have two tensors:

a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])

I’d like to obtain a matrix of shape a x b, where each value is obtained using a custom function. For the sake of the example, let’s say f(a,b) = a * b + 2.

So, my output would be:

f(1,4) | f(1,5) | f(1,6)
f(2,4) | f(2,5) | f(2,6)
f(3,4) | f(3,5) | f(3,6)

How do I do this without using loops?

please dont tag specific people.

without loops, this is best done by first replicating tensors and then doing the operations:

a.view(-1, 1).repeat(1, 3) * b.view(1, -1).repeat(3, 1) + 2

Thanks for your response. I was able to do it using Numpy’s frompyfunc. I don’t suppose there is a Pytorch equivalent of that?

And thanks for the feedback, I won’t tag specific people from now on.

there are but they are not going to be efficient.

There’s an experimental package called torchdim that does this really elegantly and efficiently. For it, however, you currently need the PyTorch nightly (not yet supported with a released PyTorch): GitHub - facebookresearch/torchdim: Named tensors with first-class dimensions for PyTorch

1 Like