Add example of how to run on PyMC3 models about blackjax HOT 24 CLOSED

from aesara.link.jax.dispatch import jax_funcify
from aesara.graph.fg import FunctionGraph

seed = jax.random.PRNGKey(1234)
chains = 1


fgraph = FunctionGraph(model.free_RVs, [model.logpt])
fns = jax_funcify(fgraph)
logp_fn_jax = fns[0]

rv_names = [rv.name for rv in model.free_RVs]
init_state = [model.test_point[rv_name] for rv_name in rv_names]
init_state_batched = jax.tree_map(lambda x: np.repeat(x[None, ...], chains, axis=0), init_state)

fns = jax_funcify(fgraph)
logp_fn_jax = fns[0]


MissingInputError                         Traceback (most recent call last)

<ipython-input-9-6e87cadc4860> in <module>()
      6 
      7 
----> 8 fgraph = FunctionGraph(model.free_RVs, [model.logpt])
      9 fns = jax_funcify(fgraph)
     10 logp_fn_jax = fns[0]

2 frames

/usr/local/lib/python3.7/dist-packages/aesara/graph/fg.py in import_node(self, apply_node, check, reason, import_missing)
    400                                 "for more information on this error."
    401                             )
--> 402                             raise MissingInputError(error_msg, variable=var)
    403 
    404         for node in new_nodes:

MissingInputError: Input 0 (TensorConstant{-inf}) of the graph (indices start from 0), used to compute InplaceDimShuffle{x}(TensorConstant{-inf}), was not provided and not given a value. Use the Aesara flag exception_verbosity='high', for more information on this error.
 
Backtrace when that variable is created:

  File "/usr/local/lib/python3.7/dist-packages/IPython/core/interactiveshell.py", line 2882, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-5-380e30b89547>", line 8, in <module>
    obs = pm.Normal("obs", mu=theta, sigma=sigma, shape=J, observed=y)
  File "/usr/local/lib/python3.7/dist-packages/pymc3/distributions/distribution.py", line 122, in __new__
    return model.Var(name, dist, data, total_size, dims=dims)
  File "/usr/local/lib/python3.7/dist-packages/pymc3/model.py", line 1182, in Var
    model=self,
  File "/usr/local/lib/python3.7/dist-packages/pymc3/model.py", line 1831, in __init__
    self.logp_sum_unscaledt = distribution.logp_sum(data)
  File "/usr/local/lib/python3.7/dist-packages/pymc3/distributions/distribution.py", line 267, in logp_sum
    return tt.sum(self.logp(*args, **kwargs))
  File "/usr/local/lib/python3.7/dist-packages/pymc3/distributions/continuous.py", line 535, in logp
    return bound((-tau * (value - mu) ** 2 + tt.log(tau / np.pi / 2.0)) / 2.0, sigma > 0)
  File "/usr/local/lib/python3.7/dist-packages/pymc3/distributions/dist_math.py", line 82, in bound
    return tt.switch(alltrue(conditions), logp, -np.inf)
`

from aesara.link.jax.dispatch import jax_funcify
from aesara.graph.fg import FunctionGraph

seed = jax.random.PRNGKey(1234)
chains = 1


fgraph = FunctionGraph(outputs= [model.logpt])
fns = jax_funcify(fgraph)
logp_fn_jax = fns[0]

rv_names = [rv.name for rv in model.free_RVs]
init_state = [model.test_point[rv_name] for rv_name in rv_names]
init_state_batched = jax.tree_map(lambda x: np.repeat(x[None, ...], chains, axis=0), init_state)

fns = jax_funcify(fgraph)
logp_fn_jax = fns[0]

WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

---------------------------------------------------------------------------

AttributeError                            Traceback (most recent call last)

<ipython-input-7-c0d3406b2bee> in <module>()
      6 
      7 
----> 8 fgraph = FunctionGraph(outputs= [model.logpt])
      9 fns = jax_funcify(fgraph)
     10 logp_fn_jax = fns[0]

3 frames

/usr/local/lib/python3.7/dist-packages/aesara/graph/fg.py in setup_node(self, node)
    212                 " the values must be tuples or lists."
    213             )
--> 214         if node.op.destroy_map and not all(
    215             isinstance(destroy, (list, tuple))
    216             for destroy in node.op.destroy_map.values()

AttributeError: 'DimShuffle' object has no attribute 'destroy_map'

---------------------------------------------------------------------------

NotImplementedError                       Traceback (most recent call last)

<ipython-input-9-c0d3406b2bee> in <module>()
      7 
      8 fgraph = FunctionGraph(outputs= [model.logpt])
----> 9 fns = jax_funcify(fgraph)
     10 logp_fn_jax = fns[0]
     11 

4 frames

/usr/local/lib/python3.7/dist-packages/aesara/link/jax/dispatch.py in jax_funcify(op, node, storage_map, **kwargs)
    141 def jax_funcify(op, node=None, storage_map=None, **kwargs):
    142     """Create a JAX compatible function from an Aesara `Op`."""
--> 143     raise NotImplementedError(f"No JAX conversion for the given `Op`: {op}")
    144 
    145 

NotImplementedError: No JAX conversion for the given `Op`: TransformedVariable