Clarification: Mapping FullyConnectedTensorProduct and Gate from e3nn to e3nn-jax

[nkzk-stan]

I know there has been some previous discussions between mapping of the FullyConnectedTensorProduct in e3nn to an equivalent form in e3nn-jax. I was seeking some further clarification on this topic as well as the use of Gate.

FullyConnectedTensorProduct in e3nn vs. e3nn-jax:
In e3nn, I employ FullyConnectedTensorProduct in the following manner:

tp = FullyConnectedTensorProduct(
          irreps_in1="irreps_in1",
          irreps_in2="irreps_in2",
          irreps_out="irreps_out",
          shared_weights=True,
          normalization='component'
      )

Subsequently, I utilize tp for various operations such as tp.weight_views(), tp.instructions, tp.weight_numel, and invoking it directly with tp(data1, data2). Could you please provide detailed guidance on achieving these operations within e3nn-jax? Specifically, how can we replicate the FullyConnectedTensorProduct functionality and the associated methods in e3nn-jax?

Gate Functionality in e3nn vs. e3nn-jax:
Regarding the Gate mechanism, I understand that e3nn-jax offers enhanced control over its behavior compared to e3nn. In e3nn, my implementation looks something like this:

Gate(
    irreps_g_scalars="irreps_g_scalars",
    act_scalars=[nn.SiLU()],
    irreps_g_gate="irreps_g_gate",
    act_gates=[torch.sigmoid],
    irreps_g_gated="irreps_g_gated"
)

I'm particularly interested in how to translate this Gate configuration to e3nn-jax, factoring in the additional control and customization capabilities it provides.

Thanks!

[nkzk-stan]

^ bumping. thank you

[mariogeiger]

FullyConnectedTensorProduct in e3nn-jax:

x1 = e3nn.normal("0e + 1o", jax.random.PRNGKey(0), ())
x2 = e3nn.normal("0e + 1o", jax.random.PRNGKey(0), ())

y = e3nn.tensor_product(x1, x2)

lin = e3nn.flax.Linear("3x0e")
w = lin.init(jax.random.PRNGKey(0), y)
lin.apply(w, y)

Gate Functionality in e3nn-jax:

input = e3nn.IrrepsArray(
    "0e + 0o + 1o + 1o + 0e + 0o",
    jnp.array([1.0, 1.0, 10.0, 0.0, 0.0, 10.0, 0.0, 0.0, 1.0, 1.0]),
)
# input = [se, so, v1, v2, ge, go]

output = e3nn.gate(
    input,
    even_act=jnp.tanh,
    odd_act=jnp.tanh,
    even_gate_act=jnp.tanh,
    odd_gate_act=jnp.tanh,
)
# output = even_act(se), odd_act(so), even_gate_act(ge) * v1, odd_gate_act(go) * v2

documentation