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!