-
Notifications
You must be signed in to change notification settings - Fork 13
Open
Labels
Description
In
s2fft/s2fft/recursions/trapani.py
Lines 776 to 792 in ea11832
| def _arg_checks(dl: np.ndarray, L: int, el: int): | |
| r""" | |
| Check arguments of Trapani functions. | |
| Args: | |
| dl (np.ndarray): Wigner-d plane of which to check shape. | |
| L (int): Harmonic band-limit. | |
| el (int): Spherical harmonic degree :math:`\ell`. | |
| """ | |
| # assert 0 < el < L | |
| # assert dl.shape[0] == dl.shape[1] == 2 * L - 1 | |
| # if L > 1024: | |
| # logs.warning_log("Trapani recursion may not be stable for L > 1024") |
the entire body of the function is commented out. I suspect this may be something to do with the assert statements not working correctly within a JAX JIT compiled function due to needing to know the concrete shapes of arrays given the commenting out was done in the changes in 3b514f1 adding JAX implementation by @jasonmcewen. We should either remove the function if no longer required or update to make it work if the issue is with JAX JIT transform.