given a matrix
z of shape N x P (samples x features), I’d like to construct a matrix with all possible interaction terms, i.e. the final result should be a matrix of size N x (P/2 * (P+1)), e.g., for a
z with 3 feature columns:
[z1, z2, z3, z1z2, z1z3, z2z3].
I tried the following:
def apply_along_axis(function, x, axis=1): return torch.stack([function(x_i) for i, x_i in enumerate(torch.unbind(x, dim=axis), 0)], dim=axis) def f(x): y=torch.triu(torch.outer(x,x), diagonal=1).flatten() return torch.cat((x, y[y!=0])) z = torch.rand(10, 3) zxz = apply_along_axis(f, z, 0) # shape: torch.Size([10, 6])
This works like a charme. I compute the outer-product (row wise), and concatenate the upper triangular part of the resulting matrix (without the diagonal) to the input row.
However, when I increase the number of particles in my
pyro.infer.Trace_ELBO instance, I observe another dimension on the left of my matrix
z, which breaks the code (since the
apply_along_axis does not pass the correct tensor into
f anymore). Any ideas, how I can construct this interaction matrix independent of the number of particles?
Thanks in advance!
PS: I also use
vectorize_particles=True which works fine in all other parts of the code (but in those cases, I am usually in a plate statement).
PPS: For interested readers, there is also https://pyro.ai/examples/sparse_regression.html