FullyConnectedTensorProduct Feature Discrepancy in JAX vs. Torch #44

[tai-dang11]

To further clarify my question regarding FullyConnectedTensorProduct: In the PyTorch version, the weight parameter can accept both a single tensor of shape (..., self.weight_numel) or a list of tensors with shapes weight_shape / (...) + weight_shape when internal_weights is set to False. In my experience with the JAX version, it seems to primarily work with the list of tensors format. Is there support, or is there a planned feature, for also using a single tensor as a weight in the JAX version, similar to the PyTorch version?

Thank you again for your assistance!

[mariogeiger]

Here is an example on how to do that in e3nn-jax

def model(x):
    y = e3nn.tensor_product(x, x)
    lin = e3nn.FunctionalLinear(y.irreps, "0e + 1o")
    w = e3nn.haiku.MultiLayerPerceptron(
        [64, 64, lin.num_weights], jax.nn.silu, output_activation=False
    )(y.filter(keep="0e").array)
    return lin(w, y)


model = hk.without_apply_rng(hk.transform(model))
x = e3nn.normal("2x0e + 3x1o + 3x1e", jax.random.PRNGKey(0))

params = model.init(jax.random.PRNGKey(1), x)
y = model.apply(params, x)

[tai-dang11]

Hello,

Thank you for your previous insights regarding the equivalence of FullyConnectedTensorProduct to tensor_product followed by a Linear layer. I've attempted to employ this approach using Torch and it seems to function effectively. However, I'm encountering some discrepancies when trying to replicate this in Jax, specifically regarding the Linear layer in Flax.

While the tensor_product operations seem consistent between both Torch and Jax, I'm observing different weights and total parameters in the Linear layer between the two libraries, as illustrated below.

I’ve provided input data and additional context in this shared drive for further reference.

I'd appreciate it if you could help me understand if there's something I might be overlooking in the Jax/Flax implementation, or if there are inherent differences between the two that might account for these discrepancies.

Thank you for your time and assistance!

[mariogeiger]

Be careful feeding weights to e3nn.Linear has a different effect than using e3nn.FunctionalLinear.

Look again carefully at the example I gave.

[tai-dang11]

Thank you so much. I got it now