Import tree_map from jax.tree_util (#56)

Compatibility with jax 0.8.0

Author: lauradriscoll

sbijax/_src/abc/smc_abc.py

@@ -7,7 +7,7 @@
 from jax import numpy as jnp
 from jax import random as jr
 from jax import scipy as jsp
-from jax import tree_map
+from jax.tree_util import tree_map
 from jax._src.flatten_util import ravel_pytree
 from tensorflow_probability.substrates.jax import distributions as tfd
 from tqdm import tqdm