Last readout for AtomicDipolesMACE

[JoshRackers]

I am wondering if there is a problem with the last readout for the AtomicDipolesMACE model. If I look at the model printout, the readouts look like this:

  (readouts): ModuleList(
    (0): LinearGenericReadoutBlock(
      (linear): Linear(10x0e+10x1o+10x2e+10x3o -> 1x0e+1x1o | 20 weights)
    )
    (1): LinearGenericReadoutBlock(
      (linear): Linear(10x0e+10x1o+10x2e+10x3o -> 1x0e+1x1o | 20 weights)
    )
    (2): NonLinearGenericReadoutBlock(
      (equivariant_nonlin): Gate (16x0e -> 16x0e)
      (linear_1): Linear(10x0e+10x1o+10x2e+10x3o -> 16x0e | 160 weights)
      (linear_2): Linear(16x0e -> 1x0e+1x1o | 16 weights)
    )
  )

The problem is the last linear_2. The model happily outputs both scalars and vectors, but all of the vectors are zero. I think this is because there is not a way to get non-zero vectors from a linear layer if the input is just scalars? Am I missing somthing?

[ilyes319]

Hey @JoshRackers,

What branch/version are you using because I don't recognize the name of the functions. Normally this should call the module NonLinearDipoleReadoutBlock and not NonLinearGenericReadoutBlock.

[JoshRackers]

Sorry, I made a slight modification on the develop branch to allow readout of arbitrary irreps_out. It seems that the answer is to change MLP_irreps to contain irreps up to the max in irreps_out. Thanks!

[davkovacs]

@JoshRackers Feel free to submit a pull request if you think that your enhancement works and others could also benefit from it.

[JoshRackers]

Sure thing! I'll submit a PR once I'm sure it's working correctly