def rec_conv_to_rv(v, replacements, model, rand_state=None): """Recursively convert a PyMC3 random variable to a Theano graph.""" if v in replacements: return walk(v, replacements) elif v.name and pm.util.is_transformed_name(v.name): untrans_name = pm.util.get_untransformed_name(v.name) v_untrans = getattr(model, untrans_name) rv_new = rec_conv_to_rv(v_untrans, replacements, model, rand_state=rand_state) replacements[v] = rv_new return rv_new elif hasattr(v, "distribution"): rv = pymc3_var_to_rv(v, rand_state=rand_state) rv_ins = [] for i in tt_inputs([rv]): i_rv = rec_conv_to_rv(i, replacements, model, rand_state=rand_state) if i_rv is not None: replacements[i] = i_rv rv_ins.append(i_rv) else: rv_ins.append(i) _ = replace_input_nodes(rv_ins, [rv], memo=replacements, clone_inputs=False) rv_new = walk(rv, replacements) replacements[v] = rv_new return rv_new else: return None
def model_graph(pymc_model, output_vars=None, rand_state=None, attach_memo=True): """Convert a PyMC3 model into a Theano `FunctionGraph`. Parameters ---------- pymc_model: `Model` A PyMC3 model object. output_vars: list (optional) Variables to use as `FunctionGraph` outputs. If not specified, the model's observed random variables are used. rand_state: Numpy rng (optional) When converting to `RandomVariable`s, use this random state object. attach_memo: boolean (optional) Add a property to the returned `FunctionGraph` name `memo` that contains the mappings between PyMC and `RandomVariable` terms. Results ------- out: `FunctionGraph` """ model = pm.modelcontext(pymc_model) replacements = {} if output_vars is None: output_vars = list(model.observed_RVs) if rand_state is None: rand_state = theano.shared(np.random.RandomState()) replacements = {} # First pass... for i, o in enumerate(output_vars): _ = rec_conv_to_rv(o, replacements, model, rand_state=rand_state) output_vars[i] = walk(o, replacements) output_vars = [walk(o, replacements) for o in output_vars] fg_features = [tt.opt.ShapeFeature()] model_fg = FunctionGraph( [i for i in tt_inputs(output_vars) if not isinstance(i, tt.Constant)], output_vars, clone=True, memo=replacements, features=fg_features, ) if attach_memo: model_fg.memo = replacements return model_fg
def constant_neq_goal(s): lvar_val = walk(lvar, s) if isinstance(lvar_val, (tt.Constant, TheanoMetaConstant)): if lvar_val.data != val: yield s else: yield s
def _reify_ConstrainedState(u, S): u_res = walk(u, S.data) if u_res is u: yield ConstrainedVar(u_res, S) else: yield _reify(u_res, S)
def _goal(s): lvar_val = walk(lvar, s) if isinstance(lvar_val, (tt.Constant, MetaConstant)): data = lvar_val.data if ((isinstance(val, np.ndarray) and not np.array_equal(data, val)) or not all(np.atleast_1d(data) == val)): yield s else: yield s
def reify_stack(u, s): u_ = walk(u, s) if u_ is not u: return reify_stack(u_, s) if isinstance(u_, (tuple, list)): return type(u_)(reify_stack(i_u, s) for i_u in u_) return u_
def unify_stack(u, v, s): u = walk(u, s) v = walk(v, s) if u == v: return s if isvar(u): return assoc(s, u, v) if isvar(v): return assoc(s, v, u) if isinstance(u, (tuple, list)) and type(u) == type(v): for i_u, i_v in zip(u, v): s = unify_stack(i_u, i_v, s) if s is False: return s return s return False
def _unify_ConstrainedVar_object(u, v, s): u_w = walk(u, s) if isvar(v): v_w = walk(v, s) else: v_w = v if u_w == v_w: yield s elif isvar(u_w): if (not isvar(v_w) and isinstance(u_w, ConstrainedVar) and not u_w.constraint(eval_if_etuple(v_w))): yield False return yield assoc(s, u_w, v_w) elif isvar(v_w): if (not isvar(u_w) and isinstance(v_w, ConstrainedVar) and not v_w.constraint(eval_if_etuple(u_w))): yield False return yield assoc(s, v_w, u_w) else: yield _unify(u_w, v_w, s)
def assocunify(u, v, s, eq=core.eq, n=None): """ Associative Unification See Also: eq_assoccomm """ uop, uargs = op_args(u) vop, vargs = op_args(v) if not uop and not vop: res = unify(u, v, s) if res is not False: # return Stream((res, )) return (res, ) # TODO: iterate through all possibilities if uop and vop: s = unify(uop, vop, s) if s is False: return ().__iter__() op = walk(uop, s) sm, lg = (uargs, vargs) if len(uargs) <= len(vargs) else (vargs, uargs) ops = assocsized(op, lg, len(sm)) goal = condeseq([(eq, a, b) for a, b, in zip(sm, lg2)] for lg2 in ops) return goaleval(goal)(s) if uop: op, tail = uop, uargs b = v if vop: op, tail = vop, vargs b = u ns = [n] if n else range(2, len(tail) + 1) knowns = (build(op, x) for n in ns for x in assocsized(op, tail, n)) goal = condeseq([(core.eq, b, k)] for k in knowns) return goaleval(goal)(s)
def assocunify(u, v, s, eq=core.eq, n=None): """ Associative Unification See Also: eq_assoccomm """ uop, uargs = op_args(u) vop, vargs = op_args(v) if not uop and not vop: res = unify(u, v, s) if res is not False: return (res, ) # TODO: iterate through all possibilities if uop and vop: s = unify(uop, vop, s) if s is False: return ().__iter__() op = walk(uop, s) sm, lg = (uargs, vargs) if len(uargs) <= len(vargs) else (vargs, uargs) ops = assocsized(op, lg, len(sm)) goal = condeseq([(eq, a, b) for a, b, in zip(sm, lg2)] for lg2 in ops) return goaleval(goal)(s) if uop: op, tail = uop, uargs b = v if vop: op, tail = vop, vargs b = u ns = [n] if n else range(2, len(tail) + 1) knowns = (build(op, x) for n in ns for x in assocsized(op, tail, n)) goal = condeseq([(core.eq, b, k)] for k in knowns) return goaleval(goal)(s)
def test_pymc_normal_model(): """Conduct a more in-depth test of PyMC3/Theano conversions for a specific model.""" tt.config.compute_test_value = 'ignore' mu_X = tt.dscalar('mu_X') sd_X = tt.dscalar('sd_X') mu_Y = tt.dscalar('mu_Y') mu_X.tag.test_value = np.array(0., dtype=tt.config.floatX) sd_X.tag.test_value = np.array(1., dtype=tt.config.floatX) mu_Y.tag.test_value = np.array(1., dtype=tt.config.floatX) # We need something that uses transforms... with pm.Model() as model: X_rv = pm.Normal('X_rv', mu_X, sd=sd_X) S_rv = pm.HalfCauchy('S_rv', beta=np.array(0.5, dtype=tt.config.floatX)) Y_rv = pm.Normal('Y_rv', X_rv * S_rv, sd=S_rv) Z_rv = pm.Normal('Z_rv', X_rv + Y_rv, sd=sd_X, observed=10.) fgraph = model_graph(model, output_vars=[Z_rv]) Z_rv_tt = canonicalize(fgraph, return_graph=False) # This will break comparison if we don't reuse it rng = Z_rv_tt.owner.inputs[1].owner.inputs[-1] mu_X_ = mt.dscalar('mu_X') sd_X_ = mt.dscalar('sd_X') tt.config.compute_test_value = 'ignore' X_rv_ = mt.NormalRV(mu_X_, sd_X_, None, rng, name='X_rv') S_rv_ = mt.HalfCauchyRV(np.array(0., dtype=tt.config.floatX), np.array(0.5, dtype=tt.config.floatX), None, rng, name='S_rv') Y_rv_ = mt.NormalRV(mt.mul(X_rv_, S_rv_), S_rv_, None, rng, name='Y_rv') Z_rv_ = mt.NormalRV(mt.add(X_rv_, Y_rv_), sd_X, None, rng, name='Z_rv') obs_ = mt(Z_rv.observations) Z_rv_obs_ = mt.observed(obs_, Z_rv_) Z_rv_meta = mt(canonicalize(Z_rv_obs_.reify(), return_graph=False)) assert mt(Z_rv_tt) == Z_rv_meta # Now, let's try that with multiple outputs. fgraph.disown() fgraph = model_graph(model, output_vars=[Y_rv, Z_rv]) assert len(fgraph.variables) == 25 Y_new_rv = walk(Y_rv, fgraph.memo) S_new_rv = walk(S_rv, fgraph.memo) X_new_rv = walk(X_rv, fgraph.memo) Z_new_rv = walk(Z_rv, fgraph.memo) # Make sure our new vars are actually in the graph and where # they should be. assert Y_new_rv == fgraph.outputs[0] assert Z_new_rv == fgraph.outputs[1] assert X_new_rv in fgraph.variables assert S_new_rv in fgraph.variables assert isinstance(Z_new_rv.owner.op, Observed) # Let's only look at the variables involved in the `Z_rv` subgraph. Z_vars = theano.gof.graph.variables(theano.gof.graph.inputs([Z_new_rv]), [Z_new_rv]) # Let's filter for only the `RandomVariables` with names. Z_vars_count = Counter([ n.name for n in Z_vars if n.name and n.owner and isinstance(n.owner.op, RandomVariable) ]) # Each new RV should be present and only occur once. assert Y_new_rv.name in Z_vars_count.keys() assert X_new_rv.name in Z_vars_count.keys() assert Z_new_rv.owner.inputs[1].name in Z_vars_count.keys() assert all(v == 1 for v in Z_vars_count.values())
def replace_input_nodes(inputs, outputs, replacements=None, memo=None, clone_inputs=True): """Recreate a graph, replacing input variables according to a given map. This is helpful if you want to replace the variable dependencies of an existing variable according to a `clone_get_equiv` map and/or replacement variables that already exist within a `FunctionGraph`. The latter is especially annoying, because you can't simply make a `FunctionGraph` for the variable to be adjusted and then use that to perform the replacement; if the variables to be replaced are already in a `FunctionGraph` any such replacement will err-out saying "...these variables are already owned by another graph..." Parameters ---------- inputs: list List of input nodes. outputs: list List of output nodes. Everything between `inputs` and these `outputs` is the graph under consideration. replacements: dict (optional) A dictionary mapping existing nodes to their new ones. These values in this map will be used instead of newly generated clones. This dict is not altered. memo: dict (optional) A dictionary to update with the initial `replacements` and maps from any old-to-new nodes arising from an actual replacement. It serves the same role as `replacements`, but it is updated as elements are cloned. clone_inputs: bool (optional) If enabled, clone all the input nodes that aren't mapped in `replacements`. These cloned nodes are mapped in `memo`, as well. Results ------- out: memo """ if memo is None: memo = {} if replacements is not None: memo.update(replacements) for apply in io_toposort(inputs, outputs): walked_inputs = [] for i in apply.inputs: if clone_inputs: # TODO: What if all the inputs are in the memo? walked_inputs.append(memo.setdefault(i, i.clone())) else: walked_inputs.append(walk(i, memo)) if any(w != i for w, i in zip(apply.inputs, walked_inputs)): new_apply = apply.clone_with_new_inputs(walked_inputs) memo.setdefault(apply, new_apply) for output, new_output in zip(apply.outputs, new_apply.outputs): memo.setdefault(output, new_output) return memo