Update custom_ops.py #315
Update custom_ops.py #315
Conversation
Small compatibility change which disables jitting on the s2fft side, in turn enables higher level jitting in s2ai.
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #315 +/- ##
==========================================
- Coverage 96.51% 96.50% -0.01%
==========================================
Files 32 32
Lines 3439 3436 -3
==========================================
- Hits 3319 3316 -3
Misses 120 120 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
@kmulderdas - is having these functions jitted blocking you jitting a function which calls them in s2ai? Generally jax.jit wrapped functions should be able to be arbitrarily nested and there is some suggestion it can help with compilation times when a lower-level function is reused multiple times in a higher level function. I suspect removing the jit decorators on the utility functions here won't have any major performance implications but it would be good to understand why its causing issues as we widely apply jax.jit to other functions in s2fft which are likely to be contained in higher-level jitted functions.
| @@ -86,7 +86,7 @@ def wigner_subset_to_s2( | |||
| return np.fft.ifft(x, axis=-2, norm="forward") | |||
|
|
|||
|
|
|||
| @partial(jit, static_argnums=(3, 4)) | |||
| # @partial(jit, static_argnums=(3, 4)) | |||
There was a problem hiding this comment.
| # @partial(jit, static_argnums=(3, 4)) |
We generally shouldn't comment out code as we can always recover snippets from git history - this is likely to be what is causing the linting failures.
| @@ -209,7 +209,7 @@ def so3_to_wigner_subset( | |||
| return s2_to_wigner_subset(x, spins, DW, L, sampling) | |||
|
|
|||
|
|
|||
| @partial(jit, static_argnums=(3, 4, 5)) | |||
| # @partial(jit, static_argnums=(3, 4, 5)) | |||
There was a problem hiding this comment.
| # @partial(jit, static_argnums=(3, 4, 5)) |
| @@ -338,7 +338,7 @@ def s2_to_wigner_subset( | |||
| return x * (2.0 * np.pi) ** 2 | |||
|
|
|||
|
|
|||
| @partial(jit, static_argnums=(3, 4)) | |||
| # @partial(jit, static_argnums=(3, 4)) | |||
There was a problem hiding this comment.
| # @partial(jit, static_argnums=(3, 4)) |
Small compatibility change which disables jitting on the s2fft side, in turn enables higher level jitting in s2ai.