Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Conversation

@SimonWagnerHD
Copy link
Contributor

Remove the unnecessary for loop over multiplicities when using the function D_from_angles in the Irreps class.

Description

Line 743 in _irreps.py direct_sum(*[ir.D_from_angles(alpha, beta, gamma, k) for mul, ir in self for _ in range(mul)]) performs a for loop over multiplicities which in each iteration calculates ir.D_from_angles even though the result is the same each time. The suggested change calculates D_from_angles only once per irrep and repeats the resulting block according to the multiplicity.

Motivation and Context

The change significantly increases the speed of the function especially in scenarios, where the number of multiplicities is high.

How Has This Been Tested?

The changes were tested by comparing the results of old and new implementation on an irreps object Irreps("128x0e + 64x1o + 16x2e + 4x3o + 1x4e") for random tensors of angles and k and asserting that the resulting tensors are the same.

Checklist:

  • I have read the CONTRIBUTING document.
  • My code follows the code style of this project.
  • I have updated the documentation (if relevant).
  • I have added tests that cover my changes (if relevant).
  • The modified code is cuda compatible (github tests don't test cuda) (if relevant).
  • I have updated the CHANGELOG.

@mitkotak
Copy link
Member

Hey thanks for the PR ! Can we have a comparison plot just to note down the speedup for future reference ?

@SimonWagnerHD
Copy link
Contributor Author

I tested the runtime of new and old version on a NVIDIA GeForce RTX 2080 Ti. I used the same random tensors of size 1000 for the angles while changing the irreps according to "{m}x0e+{m}x1o+{m}x2e" with m in [1, 2, 4, 8, 16, 32, 64].

gpu_runtime

@mitkotak mitkotak merged commit f6f909c into e3nn:main Oct 7, 2025
3 checks passed
@mitkotak
Copy link
Member

mitkotak commented Oct 7, 2025

Thanks for the PR ! Should be available in the latest release 0.5.8

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants