Possibility of incrementing torch-harmonics and imageio versions

[pzharrington]

I'm wondering if it would be possible to relax these two constraints in requirements.txt:

imageio<=2.27.0
torch-harmonics==0.7.4

They are currently blockers for us integrating the model into earth2studio, as they're causing conflicts with other items in the package that uv is unable to resolve around.

On torch-harmonics, I see you needed a local patch here, would be curious to hear some more details on the problem there and if there could be a path towards resolution by upstreaming a fix into torch-harmonics? A minimum of v0.8.0 for torch-harmonics would be ideal for us on the earth2studio side.

[oliverwm1]

On torch-harmonics, I see you needed a local patch here, would be curious to hear some more details on the problem there and if there could be a path towards resolution by upstreaming a fix into torch-harmonics? A minimum of v0.8.0 for torch-harmonics would be ideal for us on the earth2studio side.

Without this patch, we get a runtime error when optimizing over multiple forward steps with SFNO. Specifically:

2025-04-21T21:03:18.828Z [rank5]: RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.FloatTensor [181, 180, 180]] is at version 3; expected version 2 instead. Hint: enable anomaly detection to find the 
operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True).

We spent a while with @bonevbs try to figure out the root cause but were unsuccessful. @mcgibbon may know more details.

Regardless, yes, I expect we should be able to bump torch-harmonics to 0.8.0, I'm just a little nervous about not setting a max for the version given our use of private API.

imageio<=2.27.0

We had some issue with wandb video generation with a more recent imageio version but that was a couple years ago so I expect we may be able to relax this now. Will try.

[mcgibbon]

This is an issue which I have tried to fix (i.e. get us directly using torch-harmonics) multiple times without success. Each time after debugging for about a day I figure out what the issue is, and I can never remember the next time.

It would "fix" the issue for torch-harmonics to do what our patch does, but what we do is the "wrong" thing - by not registering the precomputed tensors, they don't get transferred between devices when the model is transferred, which would be unexpected by the average user. But we find this necessary to avoid the RuntimeError.

An easier option may be for us to vendorize the torch-harmonics code we're using, the way we've done with cuhpx. This would remove our direct dependency on torch-harmonics entirely, making its version unconstrained. While that means we have to update our repo to get any new changes in torch-harmonics, that's already the case given our situation with this RuntimeError.

[pzharrington]

I see. While I agree it would be ideal to not have to patch torch-harmonics within fme, I think it is workable for earth2studio if you keep the patch but just increment the versions to something like

imageio>=2.33
torch-harmonics>=0.8.0

If you're wary about the unrestricted upper end of torch-harmonics version, I think we could also work with pinning it to 0.8.0.

Another idea -- since earth2studio would only be operating the model in inference mode, the wandb and multi-step gradient issues are not really relevant functionality anyways, so you could consider adding an optional pip install dependency group like fme[inference] that does not restrict imageio or torch-harmonics (or otherwise breaking the deps into optional groups that separate out training vs. inference requirements). I think that could also work on our end.

[ankurmahesh]

(Apologies for chiming in here.) I have been using ACE with the a near-recent torch-harmonics (v0.7.6). As mentioned above, I also needed to have the patch in sht_fix.py with the latest torch to avoid the error.

With a more recent torch-harmonics, there are a few small changes necessary to sht_fix.py. These are cosmetic, not fundamental, due to some small API changes (returning variables as torch Tensors instead of numpy arrays, appending a _ to the names of imported methods). But hopefully this saves you a bit of time with the upgrade. I have attached the modified sht_fix.py to this file, since I think I saw external PRs aren't currently accepted. But if there would be a more convenient way to share these changes let me know!

sht_fix.py

[spencerkclark]

Thanks @ankurmahesh—we arrived at a similar set of changes independently. I ported the changes I made in our internal repo to this one in #49, which pins torch-harmonics to version 0.8.0 and sets a minimum imageio version of 2.28.1.

@pzharrington let us know if that looks workable for you.

[pzharrington]

Looks great!

[brianhenn]

@pzharrington you can now install the version of fme with the updated dependencies from PyPI: https://pypi.org/project/fme/2025.11.0/