🐛[BUG]: Incorrect shift parameter in torch.roll inside forward()

[mhr123456789]

Version

1.3.0

On which installation method(s) does this occur?

No response

Describe the issue

Hi, I found a bug in the forward method.

In line 299, the code currently uses:

shifted_x = torch.roll(x, shifts=(-shift_lat, -shift_lat), dims=(1, 2))
It seems the second shift argument should be shift_lon instead of shift_lat:shifted_x = torch.roll(x, shifts=(-shift_lat, -shift_lon), dims=(1, 2))
This will cause issues when shift_lat != shift_lon (e.g. non-square windows).


### Minimum reproducible example

```shell

Relevant log output

Environment details

[coreyjadams]

Hi @mhr123456789 can you elaborate on where this is happening? A small reproducer would be very valuable.

[peterdsharpe]

Looks like this is coming from physicsnemo.nn.module.transformer_layers.py, line 295 (currently).

From initial glance, this looks like a valid bug, and it's actually present in two places:

The forward roll on line 295 uses -shift_lat for both dims:

            shifted_x = torch.roll(x, shifts=(-shift_lat, -shift_lat), dims=(1, 2))

But the inverse roll on line 322 correctly uses shift_lat and shift_lon:

            x = torch.roll(shifted_x, shifts=(shift_lat, shift_lon), dims=(1, 2))

Same pattern in the 3D case circa line 144 of the same file - looks incorrect in the forward roll, but correct in the backward roll.

That said, these are mostly used in weather related cases, so perhaps @pzharrington @NickGeneva can verify that this is indeed a bug?

[pzharrington]

Based on the comments in this file it looks like the pangu-style attention modules were based on WeatherLearn. It looks like this is indeed a bug that propagated from the WeatherLearn source when this model was integrated. It was fixed in a later commit there so we should do the same here.