Add JAX-based find_MAP

Author: jessegrabowski

pymc_experimental/__init__.py

@@ -16,6 +16,7 @@
 from pymc_experimental import gp, statespace, utils
 from pymc_experimental.distributions import *
 from pymc_experimental.inference.fit import fit
+from pymc_experimental.inference.jax_find_map import find_MAP
 from pymc_experimental.model.marginal.marginal_model import MarginalModel, marginalize
 from pymc_experimental.model.model_api import as_model
 from pymc_experimental.version import __version__