Use backwards compatible tree_map import

Author: matt-graham

s2fft/utils/torch_wrapper.py

@@ -37,7 +37,7 @@
 
 import jax
 import jax.dlpack
-from jax.tree import map as tree_map
+from jax.tree_util import tree_map
 
 try:
     import torch