From 2f1f61f797386742cb0d9ea788bc3c7c522f77b0 Mon Sep 17 00:00:00 2001 From: "Brandon T. Willard" Date: Mon, 4 Mar 2019 18:38:03 -0600 Subject: [PATCH 1/3] Start separating generic and Theano-specific meta types and functionality --- requirements.txt | 4 + symbolic_pymc/__init__.py | 282 ------------- symbolic_pymc/meta.py | 479 +++-------------------- symbolic_pymc/relations/__init__.py | 4 +- symbolic_pymc/relations/conjugates.py | 8 +- symbolic_pymc/relations/distributions.py | 2 +- symbolic_pymc/relations/linalg.py | 2 +- symbolic_pymc/tensorflow/__init__.py | 3 + symbolic_pymc/tensorflow/meta.py | 41 ++ symbolic_pymc/tensorflow/unify.py | 39 ++ symbolic_pymc/theano/__init__.py | 0 symbolic_pymc/theano/meta.py | 469 ++++++++++++++++++++++ symbolic_pymc/{rv.py => theano/ops.py} | 0 symbolic_pymc/{ => theano}/opt.py | 0 symbolic_pymc/{ => theano}/printing.py | 4 +- symbolic_pymc/{ => theano}/pymc3.py | 10 +- symbolic_pymc/theano/random_variables.py | 280 +++++++++++++ symbolic_pymc/theano/unify.py | 47 +++ symbolic_pymc/theano/utils.py | 147 +++++++ symbolic_pymc/unify.py | 41 +- symbolic_pymc/utils.py | 168 -------- tests/__init__.py | 13 - tests/test_relations.py | 82 ---- tests/theano/__init__.py | 13 + tests/{ => theano}/test_conjugates.py | 6 +- tests/{ => theano}/test_kanren.py | 6 +- tests/{ => theano}/test_linalg.py | 6 +- tests/{ => theano}/test_meta.py | 16 +- tests/{ => theano}/test_printing.py | 4 +- tests/{ => theano}/test_pymc3.py | 34 +- tests/{ => theano}/test_rv.py | 2 +- tests/{ => theano}/test_unify.py | 9 +- 32 files changed, 1148 insertions(+), 1073 deletions(-) create mode 100644 symbolic_pymc/tensorflow/__init__.py create mode 100644 symbolic_pymc/tensorflow/meta.py create mode 100644 symbolic_pymc/tensorflow/unify.py create mode 100644 symbolic_pymc/theano/__init__.py create mode 100644 symbolic_pymc/theano/meta.py rename symbolic_pymc/{rv.py => theano/ops.py} (100%) rename symbolic_pymc/{ => theano}/opt.py (100%) rename symbolic_pymc/{ => theano}/printing.py (99%) rename symbolic_pymc/{ => theano}/pymc3.py (97%) create mode 100644 symbolic_pymc/theano/random_variables.py create mode 100644 symbolic_pymc/theano/unify.py create mode 100644 symbolic_pymc/theano/utils.py delete mode 100644 tests/test_relations.py create mode 100644 tests/theano/__init__.py rename tests/{ => theano}/test_conjugates.py (91%) rename tests/{ => theano}/test_kanren.py (96%) rename tests/{ => theano}/test_linalg.py (96%) rename tests/{ => theano}/test_meta.py (83%) rename tests/{ => theano}/test_printing.py (93%) rename tests/{ => theano}/test_pymc3.py (85%) rename tests/{ => theano}/test_rv.py (97%) rename tests/{ => theano}/test_unify.py (96%) diff --git a/requirements.txt b/requirements.txt index cf4ad79..4292f5f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,9 @@ Theano>=1.0.4 pymc3>=3.6 +tf-estimator-nightly>=1.14.0.dev2019021101 +tf-nightly>=1.14.1.dev20190320 +tfp-nightly>=0.7.0.dev20190320 +pymc4 @ git+https://github.com/pymc-devs/pymc4.git@master#egg=pymc4-0.0.1 multipledispatch>=0.6.0 unification>=0.2.2 kanren @ git+https://github.com/pymc-devs/kanren.git@symbolic-pymc#egg=kanren-0.2.3 diff --git a/symbolic_pymc/__init__.py b/symbolic_pymc/__init__.py index daf919f..b3c42a8 100644 --- a/symbolic_pymc/__init__.py +++ b/symbolic_pymc/__init__.py @@ -1,286 +1,4 @@ -import theano -import scipy -import theano.tensor as tt - -from functools import partial - -from .rv import RandomVariable, param_supp_shape_fn - # We need this so that `multipledispatch` initialization occurs from .unify import * -from .meta import mt - __version__ = "0.0.1" - -# Continuous Numpy-generated variates -class UniformRVType(RandomVariable): - print_name = ("U", "\\operatorname{U}") - - def __init__(self): - super().__init__("uniform", theano.config.floatX, 0, [0, 0], "uniform", inplace=True) - - def make_node(self, lower, upper, size=None, rng=None, name=None): - return super().make_node(lower, upper, size=size, rng=rng, name=name) - - -UniformRV = UniformRVType() - - -class NormalRVType(RandomVariable): - print_name = ("N", "\\operatorname{N}") - - def __init__(self): - super().__init__("normal", theano.config.floatX, 0, [0, 0], "normal", inplace=True) - - def make_node(self, mu, sigma, size=None, rng=None, name=None): - return super().make_node(mu, sigma, size=size, rng=rng, name=name) - - -NormalRV = NormalRVType() - - -class HalfNormalRVType(RandomVariable): - print_name = ("N**+", "\\operatorname{N^{+}}") - - def __init__(self): - super().__init__( - "halfnormal", - theano.config.floatX, - 0, - [0, 0], - lambda rng, *args: scipy.stats.halfnorm.rvs(*args, random_state=rng), - inplace=True, - ) - - def make_node(self, mu=0.0, sigma=1.0, size=None, rng=None, name=None): - return super().make_node(mu, sigma, size=size, rng=rng, name=name) - - -HalfNormalRV = HalfNormalRVType() - - -class GammaRVType(RandomVariable): - print_name = ("Gamma", "\\operatorname{Gamma}") - - def __init__(self): - super().__init__("gamma", theano.config.floatX, 0, [0, 0], "gamma", inplace=True) - - def make_node(self, shape, scale, size=None, rng=None, name=None): - return super().make_node(shape, scale, size=size, rng=rng, name=name) - - -GammaRV = GammaRVType() - - -class ExponentialRVType(RandomVariable): - print_name = ("Exp", "\\operatorname{Exp}") - - def __init__(self): - super().__init__("exponential", theano.config.floatX, 0, [0], "exponential", inplace=True) - - def make_node(self, scale, size=None, rng=None, name=None): - return super().make_node(scale, size=size, rng=rng, name=name) - - -ExponentialRV = ExponentialRVType() - - -# One with multivariate support -class MvNormalRVType(RandomVariable): - print_name = ("N", "\\operatorname{N}") - - def __init__(self): - super().__init__( - "multivariate_normal", - theano.config.floatX, - 1, - [1, 2], - "multivariate_normal", - inplace=True, - ) - - def make_node(self, mean, cov, size=None, rng=None, name=None): - return super().make_node(mean, cov, size=size, rng=rng, name=name) - - -MvNormalRV = MvNormalRVType() - - -class DirichletRVType(RandomVariable): - print_name = ("Dir", "\\operatorname{Dir}") - - def __init__(self): - super().__init__("dirichlet", theano.config.floatX, 1, [1], "dirichlet", inplace=True) - - def make_node(self, alpha, size=None, rng=None, name=None): - return super().make_node(alpha, size=size, rng=rng, name=name) - - -DirichletRV = DirichletRVType() - - -# A discrete Numpy-generated variate -class PoissonRVType(RandomVariable): - print_name = ("Pois", "\\operatorname{Pois}") - - def __init__(self): - super().__init__("poisson", "int64", 0, [0], "poisson", inplace=True) - - def make_node(self, rate, size=None, rng=None, name=None): - return super().make_node(rate, size=size, rng=rng, name=name) - - -PoissonRV = PoissonRVType() - - -# A SciPy-generated variate -class CauchyRVType(RandomVariable): - print_name = ("C", "\\operatorname{C}") - - def __init__(self): - super().__init__( - "cauchy", - theano.config.floatX, - 0, - [0, 0], - lambda rng, *args: scipy.stats.cauchy.rvs(*args, random_state=rng), - inplace=True, - ) - - def make_node(self, loc, scale, size=None, rng=None, name=None): - return super().make_node(loc, scale, size=size, rng=rng, name=name) - - -CauchyRV = CauchyRVType() - - -class HalfCauchyRVType(RandomVariable): - print_name = ("C**+", "\\operatorname{C^{+}}") - - def __init__(self): - super().__init__( - "halfcauchy", - theano.config.floatX, - 0, - [0, 0], - lambda rng, *args: scipy.stats.halfcauchy.rvs(*args, random_state=rng), - inplace=True, - ) - - def make_node(self, loc=0.0, scale=1.0, size=None, rng=None, name=None): - return super().make_node(loc, scale, size=size, rng=rng, name=name) - - -HalfCauchyRV = HalfCauchyRVType() - - -class InvGammaRVType(RandomVariable): - print_name = ("InvGamma", "\\operatorname{Gamma^{-1}}") - - def __init__(self): - super().__init__( - "invgamma", - theano.config.floatX, - 0, - [0, 0], - lambda rng, *args: scipy.stats.invgamma.rvs(*args, random_state=rng), - inplace=True, - ) - - def make_node(self, loc, scale, size=None, rng=None, name=None): - return super().make_node(loc, scale, size=size, rng=rng, name=name) - - -InvGammaRV = InvGammaRVType() - - -class TruncExponentialRVType(RandomVariable): - print_name = ("TruncExp", "\\operatorname{Exp}") - - def __init__(self): - super().__init__( - "truncexpon", - theano.config.floatX, - 0, - [0, 0, 0], - lambda rng, *args: scipy.stats.truncexpon.rvs(*args, random_state=rng), - inplace=True, - ) - - def make_node(self, b, loc, scale, size=None, rng=None, name=None): - return super().make_node(b, loc, scale, size=size, rng=rng, name=name) - - -TruncExponentialRV = TruncExponentialRVType() - - -# Support shape is determined by the first dimension in the *second* parameter -# (i.e. the probabilities vector) -class MultinomialRVType(RandomVariable): - print_name = ("MN", "\\operatorname{MN}") - - def __init__(self): - super().__init__( - "multinomial", - "int64", - 1, - [0, 1], - "multinomial", - supp_shape_fn=partial(param_supp_shape_fn, rep_param_idx=1), - inplace=True, - ) - - def make_node(self, n, pvals, size=None, rng=None, name=None): - return super().make_node(n, pvals, size=size, rng=rng, name=name) - - -MultinomialRV = MultinomialRVType() - - -class Observed(tt.Op): - """An `Op` that represents an observed random variable. - - This `Op` establishes an observation relationship between a random - variable and a specific value. - """ - - default_output = 0 - - def __init__(self): - self.view_map = {0: [0]} - - def make_node(self, val, rv=None): - """Make an `Observed` random variable. - - Parameters - ---------- - val: Variable - The observed value. - rv: RandomVariable - The distribution from which `val` is assumed to be a sample value. - """ - val = tt.as_tensor_variable(val) - if rv: - if rv.owner and not isinstance(rv.owner.op, RandomVariable): - raise ValueError(f"`rv` must be a RandomVariable type: {rv}") - - if rv.type.convert_variable(val) is None: - raise ValueError( - ("`rv` and `val` do not have compatible types:" f" rv={rv}, val={val}") - ) - else: - rv = tt.NoneConst.clone() - - inputs = [val, rv] - - return tt.Apply(self, inputs, [val.type()]) - - def perform(self, node, inputs, out): - out[0][0] = inputs[0] - - def grad(self, inputs, outputs): - return outputs - - -observed = Observed() diff --git a/symbolic_pymc/meta.py b/symbolic_pymc/meta.py index e00ce26..451a445 100644 --- a/symbolic_pymc/meta.py +++ b/symbolic_pymc/meta.py @@ -5,16 +5,12 @@ import numpy as np -import theano -import theano.tensor as tt - from itertools import chain -from functools import partial, wraps +from functools import partial from collections.abc import Iterator -from unification import var, isvar, Var +from unification import isvar, Var -from .rv import RandomVariable from .utils import _check_eq from multipledispatch import dispatch @@ -31,15 +27,6 @@ def metatize(obj): return _metatize(obj) -@dispatch(object) -def _metatize(obj): - try: - obj = tt.as_tensor_variable(obj) - except (ValueError, tt.AsTensorError): - raise ValueError("Could not find a MetaSymbol class for {}".format(obj)) - return _metatize(obj) - - @dispatch((set, list, tuple)) def _metatize(obj): """Convert elements of an iterable to meta objects.""" @@ -136,7 +123,7 @@ def base(self): @classmethod def base_classes(cls, mro_order=True): res = tuple(c.base for c in cls.__subclasses__()) - if cls is not MetaSymbol: + if hasattr(cls, "base"): res = (cls.base,) + res sorted(res, key=lambda c: len(c.mro()), reverse=mro_order) return res @@ -268,67 +255,26 @@ def _repr_pretty_(self, p, cycle): p.pretty(obj) -@dispatch(type) -def _metatize(obj): - """Return an existing meta type/class, or create a new one.""" - cls = MetaSymbol - while True: - try: - obj_cls = next(filter(lambda t: issubclass(obj, t.base), cls.__subclasses__())) - except StopIteration: - # The current class is the best fit. - if cls.base == obj: - return cls - # This object is a subclass of the base type. - new_type = type(f"Meta{obj.__name__}", (obj_cls,), {"base": obj}) - return new_type(obj) - else: - cls = obj_cls - - @dispatch((MetaSymbol, type(None), types.FunctionType, partial, str, dict)) def _metatize(obj): return obj -class MetaType(MetaSymbol): - base = theano.Type - - def __call__(self, name=None): - if self.obj: - return metatize(self.obj(name=name)) - return metatize(self.base.Variable)(self, name) - - -class MetaRandomStateType(MetaType): - base = tt.raw_random.RandomStateType - - -class MetaTensorType(MetaType): - base = tt.TensorType - __slots__ = ["dtype", "broadcastable", "name"] - - def __init__(self, dtype, broadcastable, name, obj=None): - super().__init__(obj=obj) - self.dtype = dtype - self.broadcastable = broadcastable - self.name = name - - class MetaOp(MetaSymbol): - """A meta object that represents Theano `Op`s. - - NOTE: By default it will use `Op.make_node`'s signature to produce meta - `Apply` node inputs, so be sure to override that signature when - `Op.make_node`'s arguments aren't one-to-one with the expected `Apply` node - inputs. See `MetaOp.__call__` for more details. + """A meta object that represents a `MetaVaribale`-producing operator. Also, make sure to override `Op.out_meta_type` and make it return the expected meta variable type, if it isn't the default: `MetaTensorVariable`. + In some cases, operators hold their own inputs and outputs + (e.g. TensorFlow), and, in others, an intermediary "application" node holds + that information. This class leaves those details up to the + implementation. """ - base = tt.Op + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.op_sig = inspect.signature(self.obj.make_node) @property def obj(self): @@ -340,98 +286,14 @@ def obj(self, x): raise ValueError("Cannot reset obj in an `Op`") self._obj = x - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.op_sig = inspect.signature(self.obj.make_node) - + @abc.abstractmethod def out_meta_type(self, inputs=None): - """Return the type of meta variable this `Op` is expected to produce given the inputs. - - The default is `MetaTensorVariable` (corresponding to - `TheanoTensorVariable` outputs from the base `Op`). - - """ - return MetaTensorVariable + """Return the type of meta variable this `Op` is expected to produce given the inputs.""" + raise NotImplementedError() + @abc.abstractmethod def __call__(self, *args, ttype=None, index=None, **kwargs): - """Emulate `make_node` for this `Op` and return . - - NOTE: Meta objects will use positional arguments and non-"name" keyword - args as `Apply` node inputs. Also, if some of the `Op` constructor - arguments that end up as `Apply` node input arguments are keywords, - *use the keywords* and not their positions! - - Otherwise, if a base object can't be referenced, unknown Theano types - and index values will be fill-in with logic variables (that can also be - specified manually though the keyword arguments `ttype` and `index`). - - Parameters - ---------- - ttype: object (optional) - Value to use for an unknown Theano type. Defaults to a logic - variable. - index: object (optional) - Value to use for an unknown output index value. Defaults to a - logic variable. - - """ - name = kwargs.pop("name", None) - - # Use the `Op`'s default `make_node` arguments, if any. - op_arg_bind = self.op_sig.bind(*args, **kwargs) - op_arg_bind.apply_defaults() - op_args, op_args_unreified = _meta_reify_iter(op_arg_bind.args) - - if not op_args_unreified: - tt_out = self.obj(*op_args) - res_var = metatize(tt_out) - - if not isinstance(tt_out, (list, tuple)): - # If the name is indeterminate, we still want all the reified info, - # but we need to make sure that certain parts aren't known. - # TODO: In this case, the reified Theano object is a sort of - # "proxy" object; we should use this approach for dtype, as well. - # TODO: We should also put this kind of logic in the appropriate places - # (e.g. `MetaVariable.reify`), when possible. - if MetaSymbol.is_meta(name): - # This should also invalidate `res_var.obj`. - res_var.name = name - # Allow the base object to be unified, so that reification - # can recover the underlying object--instead of recreating - # it and sacrificing equality. - res_var.obj = var() - - elif tt_out.name != name: - tt_out.name = name - res_var.name = name - - else: - # XXX: It's not always clear how `Op.make_node` arguments map to - # `Apply` node inputs, which is one of the big problem with - # Theano's design. (More generally, it's that `Op`s don't provide - # a spec for `Apply` node inputs and outputs at all.) - - # Also, `Apply` inputs can't be `None` (they could be - # `tt.none_type_t()`, though). - res_apply = MetaApply(self, tuple(filter(lambda x: x is not None, op_arg_bind.args))) - - # TODO: Elemwise has an `output_types` method that can be - # used to infer the output type of this variable. - ttype = ttype or var() - - # Use the given index or the base `Op`'s `default_output`; - # otherwise, create a logic variable place-holder. - index = index if index is not None else getattr(self.obj, "default_output", None) - index = index if index is not None else var() - - # XXX: We don't have a higher-order meta object model, so being - # wrong about the exact type of output variable will cause - # problems. - out_meta_type = self.out_meta_type(op_args) - res_var = out_meta_type(ttype, res_apply, index, name) - res_var.obj = var() - - return res_var + raise NotImplementedError() def __eq__(self, other): # Since these have no rands/slots, we can only really compare against @@ -450,291 +312,40 @@ def __hash__(self): return hash((self.base, self.obj)) -class MetaElemwise(MetaOp): - base = tt.Elemwise - - def __call__(self, *args, ttype=None, index=None, **kwargs): - obj_nout = getattr(self.obj, "nfunc_spec", None) - obj_nout = obj_nout[-1] if obj_nout is not None else None - if obj_nout == 1 and index is None: - index = 0 - return super().__call__(*args, ttype=ttype, index=index, **kwargs) - - -class MetaDimShuffle(MetaOp): - base = tt.DimShuffle - __slots__ = ["input_broadcastable", "new_order", "inplace"] - - def __init__(self, input_broadcastable, new_order, inplace=True, obj=None): - super().__init__(obj=obj) - self.input_broadcastable = input_broadcastable - self.new_order = new_order - self.inplace = inplace - - -class MetaRandomVariable(MetaOp): - base = RandomVariable - - def __init__(self, obj=None): - super().__init__(obj=obj) - # The `name` keyword parameter isn't an `Apply` node input, so we need - # to remove it from the automatically generated signature. - self.op_sig = self.op_sig.replace(parameters=list(self.op_sig.parameters.values())[0:4]) - - -class MetaApply(MetaSymbol): - base = tt.Apply - __slots__ = ["op", "inputs"] - - def __init__(self, op, inputs, outputs=None, obj=None): - super().__init__(obj=obj) - self.op = metatize(op) - self.inputs = tuple(metatize(i) for i in inputs) - self.outputs = outputs - - def reify(self): - if self.obj and not isinstance(self.obj, Var): - return self.obj - else: - tt_op = self.op.reify() - if not self.is_meta(tt_op): - reified_rands, any_unreified = _meta_reify_iter(self.inputs) - if not any_unreified: - tt_var = tt_op(*reified_rands) - self.obj = tt_var.owner - return tt_var.owner - return self - - @property - def nin(self): - return len(self.inputs) - - @property - def nout(self): - if self.outputs is not None: - return len(self.outputs) - elif self.obj: - return len(self.obj.outputs) - # TODO: Would be cool if we could return - # a logic variable representing this. - - -class MetaVariable(MetaSymbol): - base = theano.Variable - __slots__ = ["type", "owner", "index", "name"] - - def __init__(self, type, owner, index, name, obj=None): - super().__init__(obj=obj) - self.type = metatize(type) - self.owner = metatize(owner) - self.index = index - self.name = name - - def reify(self): - if self.obj and not isinstance(self.obj, Var): - return self.obj - - if not self.owner: - return super().reify() - - # Having an `owner` causes issues (e.g. being consistent about - # other, unrelated outputs of an `Apply` node), and, in this case, - # the `Apply` node that owns this variable needs to construct it. - reified_rands, any_unreified = _meta_reify_iter(self.rands()) - tt_apply = self.owner.obj - - if tt_apply and not isvar(tt_apply): - # If the owning `Apply` reified, then one of its `outputs` - # corresponds to this variable. Our `self.index` value should - # tell us which, but, when that's not available, we can - # sometimes infer it. - if tt_apply.nout == 1: - tt_index = 0 - # Make sure we didn't have a mismatched non-meta index value. - assert isvar(self.index) or self.index is None or self.index == 0 - # Set/replace `None` or meta value - self.index = 0 - tt_var = tt_apply.outputs[tt_index] - elif not self.is_meta(self.index): - tt_var = tt_apply.outputs[self.index] - elif self.index is None: - tt_var = tt_apply.default_output() - self.index = tt_apply.outputs.index(tt_var) - else: - return self - # If our name value is not set/concrete, then use the reified - # value's. Otherwise, use ours. - if isvar(self.name) or self.name is None: - self.name = tt_var.name - else: - tt_var.name = self.name - assert tt_var is not None - self.obj = tt_var - return tt_var - - return super().reify() - - -class MetaTensorVariable(MetaVariable): - # TODO: Could extend `theano.tensor.var._tensor_py_operators`, too. - base = tt.TensorVariable - - @property - def ndim(self): - if isinstance(self.type, MetaTensorType) and isinstance( - self.type.broadastable, (list, tuple) - ): - return len(self.type.broadcastable) - # TODO: Would be cool if we could return - # a logic variable representing this. - - -class MetaConstant(MetaVariable): - base = theano.Constant - __slots__ = ["type", "data"] - - def __init__(self, type, data, name=None, obj=None): - super().__init__(type, None, None, name, obj=obj) - self.data = data - - -class MetaTensorConstant(MetaConstant): - # TODO: Could extend `theano.tensor.var._tensor_py_operators`, too. - base = tt.TensorConstant - __slots__ = ["type", "data", "name"] - - def __init__(self, type, data, name=None, obj=None): - super().__init__(type, data, name, obj=obj) - - -class MetaSharedVariable(MetaVariable): - base = tt.sharedvar.SharedVariable - __slots__ = ["name", "type", "data", "strict"] - - def __init__(self, name, type, data, strict, obj=None): - super().__init__(type, None, None, name, obj=obj) - self.data = data - self.strict = strict - - @classmethod - def _metatize(cls, obj): - res = MetaSharedVariable( - obj.name, obj.type, obj.container.data, obj.container.strict, obj=obj - ) - return res - - -class MetaTensorSharedVariable(MetaSharedVariable): - # TODO: Could extend `theano.tensor.var._tensor_py_operators`, too. - base = tt.sharedvar.TensorSharedVariable - - -class MetaScalarSharedVariable(MetaSharedVariable): - base = tt.sharedvar.ScalarSharedVariable - - -class MetaAccessor(object): - """Create an object that can be used to implicitly convert Theano functions and object into meta objects. - - Use it like a namespace/module/package object, e.g. - - >>> mt = TheanoMetaAccessor() - >>> mt.vector('a') - MetaTensorVariable(MetaTensorType('float64', (False,), None, - obj=TensorType(float64, vector)), None, None, 'a', obj=a) - - Call it as a function to perform direct conversion to a meta - object, e.g. - - >>> mt(tt.vector('a')) - MetaTensorVariable(MetaTensorType('float64', (False,), None, - obj=TensorType(float64, vector)), None, None, 'a', obj=a) - - """ - - namespaces = [tt] +def _find_meta_type(obj_type, meta_abs_type): + cls = meta_abs_type + obj_cls = None + while True: + try: + obj_cls = next(filter(lambda t: issubclass(obj_type, t.base), cls.__subclasses__())) + except StopIteration: + # The current class is the best fit. + if cls.base == obj_type: + return cls - def __init__(self, namespace=None): - if namespace is None: - import symbolic_pymc - from symbolic_pymc import meta # pylint: disable=import-self + # The abstract meta type has no subclasses that match the given + # object type. + if obj_cls is None: + return None - self.namespaces += [symbolic_pymc, meta] + # This object is a subclass of an existing meta class' base type, + # but there is no implemented meta type for this subclass, so we + # dynamically make one. + new_type = type(f"Meta{obj_type.__name__}", (obj_cls,), {"base": obj_type}) + return new_type(obj_type) else: - self.namespaces = [namespace] - - def __call__(self, x): - return metatize(x) - - def __getattr__(self, obj): - - ns_obj = next((getattr(ns, obj) for ns in self.namespaces if hasattr(ns, obj)), None) - - if ns_obj is None: - # Try caller's namespace - frame = inspect.currentframe() - f_back = frame.f_back - if f_back: - ns_obj = f_back.f_locals.get(obj, None) - if ns_obj is None: - ns_obj = f_back.f_globals.get(obj) - - if isinstance(ns_obj, (types.FunctionType, partial)): - # It's a function, so let's provide a wrapper that converts - # to-and-from theano and meta objects. - @wraps(ns_obj) - def meta_obj(*args, **kwargs): - args = [o.reify() if hasattr(o, "reify") else o for o in args] - res = ns_obj(*args, **kwargs) - return metatize(res) - - elif isinstance(ns_obj, types.ModuleType): - # It's a sub-module, so let's create another - # `MetaAccessor` and check within there. - meta_obj = MetaAccessor(namespace=ns_obj) - else: - # Hopefully, it's convertible to a meta object... - meta_obj = metatize(ns_obj) - - if isinstance(meta_obj, (MetaSymbol, MetaSymbolType, types.FunctionType)): - setattr(self, obj, meta_obj) - return getattr(self, obj) - elif isinstance(meta_obj, MetaAccessor): - setattr(self, obj, meta_obj) - return meta_obj - else: - raise AttributeError(f"Meta object for {obj} not found.") - - -mt = MetaAccessor() -mt.dot = metatize(tt.basic._dot) - - -# -# The wrapped Theano functions will only work when the meta objects -# are fully reifiable (i.e. can be turned to Theano objects), but it's -# fairly straight-forward to adjust many of those functions so that they -# work with meta objects. -# TODO: Would be nice if we could trick Theano into using meta objects, or -# a robust use of "proxy" Theano objects -# - + cls = obj_cls -def mt_zeros(shape, dtype=None): - if not isinstance(shape, (list, tuple, MetaTensorVariable, tt.TensorVariable)): - shape = [shape] - if dtype is None: - dtype = tt.config.floatX - return mt.alloc(np.array(0, dtype=dtype), *shape) +@dispatch(type) +def _metatize(obj_type): + """Return an existing meta type/class, or create a new one.""" + for meta_type in MetaSymbol.__subclasses__(): + obj_cls = _find_meta_type(obj_type, meta_type) -mt.zeros = mt_zeros + if obj_cls is not None: + return obj_cls -def mt_diag(v, k=0): - if v.ndim == 1: - return mt.AllocDiag(k)(v) - elif v.ndim is not None and v.ndim >= 2: - return mt.diagonal(v, offset=k) - else: - raise ValueError("Input must has v.ndim >= 1.") +class MetaVariable(MetaSymbol): + pass diff --git a/symbolic_pymc/relations/__init__.py b/symbolic_pymc/relations/__init__.py index 2ecae61..e616d70 100644 --- a/symbolic_pymc/relations/__init__.py +++ b/symbolic_pymc/relations/__init__.py @@ -6,7 +6,7 @@ from unification.utils import transitive_get as walk -from ..meta import MetaConstant +from ..theano.meta import TheanoMetaConstant # Hierarchical models that we recognize. @@ -29,7 +29,7 @@ def constant_neq(lvar, val): def _goal(s): lvar_val = walk(lvar, s) - if isinstance(lvar_val, (tt.Constant, MetaConstant)): + if isinstance(lvar_val, (tt.Constant, TheanoMetaConstant)): data = lvar_val.data if (isinstance(val, np.ndarray) and not np.array_equal(data, val)) or not all( np.atleast_1d(data) == val diff --git a/symbolic_pymc/relations/conjugates.py b/symbolic_pymc/relations/conjugates.py index 09f660a..8ba3e15 100644 --- a/symbolic_pymc/relations/conjugates.py +++ b/symbolic_pymc/relations/conjugates.py @@ -1,15 +1,16 @@ -from theano.tensor.nlinalg import matrix_inverse # pylint: disable=unused-import +import theano from unification import var from kanren import conde, eq from kanren.facts import fact from . import conjugate -from .. import MvNormalRV, observed # pylint: disable=unused-import from ..unify import etuple -from ..meta import mt +from ..theano.meta import mt +mt.namespaces += [theano.tensor.nlinalg] + # The prior distribution prior_dist_mt = var("prior_dist") @@ -135,7 +136,6 @@ def create_normal_wishart_goals(): # fact(conjugates, # Y_obs_mt, wishart_posterior_exprs) - pass def conjugate_posteriors(x, y): diff --git a/symbolic_pymc/relations/distributions.py b/symbolic_pymc/relations/distributions.py index 99fea02..87c1899 100644 --- a/symbolic_pymc/relations/distributions.py +++ b/symbolic_pymc/relations/distributions.py @@ -4,8 +4,8 @@ from kanren.facts import fact from . import constant_neq, concat -from ..meta import mt from ..unify import etuple +from ..theano.meta import mt from kanren.facts import Relation diff --git a/symbolic_pymc/relations/linalg.py b/symbolic_pymc/relations/linalg.py index edc74a6..f729e3a 100644 --- a/symbolic_pymc/relations/linalg.py +++ b/symbolic_pymc/relations/linalg.py @@ -11,7 +11,7 @@ from kanren.goals import not_equalo, conso from kanren.term import term, operator, arguments -from ..meta import mt +from ..theano.meta import mt from ..unify import etuple, tuple_expression, ExpressionTuple diff --git a/symbolic_pymc/tensorflow/__init__.py b/symbolic_pymc/tensorflow/__init__.py new file mode 100644 index 0000000..294e343 --- /dev/null +++ b/symbolic_pymc/tensorflow/__init__.py @@ -0,0 +1,3 @@ +from tensorflow.python.framework import ops + +ops.disable_eager_execution() diff --git a/symbolic_pymc/tensorflow/meta.py b/symbolic_pymc/tensorflow/meta.py new file mode 100644 index 0000000..e6bfe51 --- /dev/null +++ b/symbolic_pymc/tensorflow/meta.py @@ -0,0 +1,41 @@ +import tensorflow as tf + +from multipledispatch import dispatch + +from ..meta import MetaSymbol, MetaOp, MetaVariable, _metatize + + +@dispatch(object) +def _metatize(obj): + try: + obj = tf.convert_to_tensor(obj) + except TypeError: + raise ValueError("Could not find a TensorFlow MetaSymbol class for {}".format(obj)) + return _metatize(obj) + + +class TFlowMetaSymbol(MetaSymbol): + def reify(self): + # super().reify() + + # TODO: Follow `tfp.distribution.Distribution`'s lead? + # with tf.name_scope(self.name): + # pass + pass + + +class TFlowMetaOp(MetaOp, TFlowMetaSymbol): + base = tf.Operation + + def __call__(self, *args, ttype=None, index=None, **kwargs): + pass + + +class TFTensorVariable(MetaVariable, TFlowMetaSymbol): + base = tf.Tensor + __slots__ = ["op", "value_index", "dtype"] + + def __init__(self, op, value_index, dtype): + self.op = op + self.value_index = value_index + self.dtype = dtype diff --git a/symbolic_pymc/tensorflow/unify.py b/symbolic_pymc/tensorflow/unify.py new file mode 100644 index 0000000..08beaef --- /dev/null +++ b/symbolic_pymc/tensorflow/unify.py @@ -0,0 +1,39 @@ +import tensorflow as tf + +from kanren.term import term, operator, arguments + +from unification.core import reify, _unify, _reify + +from ..meta import metatize +from ..unify import ExpressionTuple, unify_MetaSymbol +from .meta import TFlowMetaSymbol + +tf_class_abstractions = tuple(c.base for c in TFlowMetaSymbol.__subclasses__()) + +_unify.add( + (TFlowMetaSymbol, tf_class_abstractions, dict), + lambda u, v, s: unify_MetaSymbol(u, metatize(v), s), +) +_unify.add( + (tf_class_abstractions, TFlowMetaSymbol, dict), + lambda u, v, s: unify_MetaSymbol(metatize(u), v, s), +) +_unify.add( + (tf_class_abstractions, tf_class_abstractions, dict), + lambda u, v, s: unify_MetaSymbol(metatize(u), metatize(v), s), +) + + +def _reify_TFlowClasses(o, s): + meta_obj = metatize(o) + return reify(meta_obj, s) + + +_reify.add((tf_class_abstractions, dict), _reify_TFlowClasses) + + +operator.add((tf.Tensor,), lambda x: operator(metatize(x))) + +arguments.add((tf.Tensor,), lambda x: arguments(metatize(x))) + +term.add((tf.Operation, ExpressionTuple), lambda op, args: term(metatize(op), args)) diff --git a/symbolic_pymc/theano/__init__.py b/symbolic_pymc/theano/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/symbolic_pymc/theano/meta.py b/symbolic_pymc/theano/meta.py new file mode 100644 index 0000000..c505ba0 --- /dev/null +++ b/symbolic_pymc/theano/meta.py @@ -0,0 +1,469 @@ +import types +import inspect + +import theano +import theano.tensor as tt + +from functools import partial, wraps + +from unification import var, isvar, Var +from multipledispatch import dispatch + +from kanren.facts import fact +from kanren.assoccomm import commutative, associative + +from .ops import RandomVariable +from ..meta import ( + MetaSymbol, + MetaSymbolType, + MetaOp, + MetaVariable, + _meta_reify_iter, + metatize, + _metatize, +) + + +@dispatch(object) +def _metatize(obj): + try: + obj = tt.as_tensor_variable(obj) + except (ValueError, tt.AsTensorError): + raise ValueError("Could not find a MetaSymbol class for {}".format(obj)) + return _metatize(obj) + + +class TheanoMetaSymbol(MetaSymbol): + pass + + +class TheanoMetaType(TheanoMetaSymbol): + base = theano.Type + + def __call__(self, name=None): + if self.obj: + return metatize(self.obj(name=name)) + return metatize(self.base.Variable)(self, name) + + +class TheanoMetaRandomStateType(TheanoMetaType): + base = tt.raw_random.RandomStateType + + +class TheanoMetaTensorType(TheanoMetaType): + base = tt.TensorType + __slots__ = ["dtype", "broadcastable", "name"] + + def __init__(self, dtype, broadcastable, name, obj=None): + super().__init__(obj=obj) + self.dtype = dtype + self.broadcastable = broadcastable + self.name = name + + +class TheanoMetaOp(MetaOp, TheanoMetaSymbol): + """A meta object that represents Theano `Op`s. + + NOTE: By default it will use `Op.make_node`'s signature to produce meta + `Apply` node inputs, so be sure to override that signature when + `Op.make_node`'s arguments aren't one-to-one with the expected `Apply` node + inputs. See `MetaOp.__call__` for more details. + + Also, make sure to override `Op.out_meta_type` and make it return the + expected meta variable type, if it isn't the default: `TheanoMetaTensorVariable`. + """ + + base = tt.Op + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def out_meta_type(self, inputs=None): + """Return the type of meta variable this `Op` is expected to produce given the inputs. + + The default is `TheanoMetaTensorVariable` (corresponding to + `TheanoTensorVariable` outputs from the base `Op`). + """ + return TheanoMetaTensorVariable + + def __call__(self, *args, ttype=None, index=None, **kwargs): + """Emulate `make_node` for this `Op` and return . + + NOTE: Meta objects will use positional arguments and non-"name" keyword + args as `Apply` node inputs. Also, if some of the `Op` constructor + arguments that end up as `Apply` node input arguments are keywords, + *use the keywords* and not their positions! + + Otherwise, if a base object can't be referenced, unknown Theano types + and index values will be fill-in with logic variables (that can also be + specified manually though the keyword arguments `ttype` and `index`). + + Parameters + ---------- + ttype: object (optional) + Value to use for an unknown Theano type. Defaults to a logic + variable. + index: object (optional) + Value to use for an unknown output index value. Defaults to a + logic variable. + + """ + name = kwargs.pop("name", None) + + # Use the `Op`'s default `make_node` arguments, if any. + op_arg_bind = self.op_sig.bind(*args, **kwargs) + op_arg_bind.apply_defaults() + op_args, op_args_unreified = _meta_reify_iter(op_arg_bind.args) + + if not op_args_unreified: + tt_out = self.obj(*op_args) + res_var = metatize(tt_out) + + if not isinstance(tt_out, (list, tuple)): + # If the name is indeterminate, we still want all the reified info, + # but we need to make sure that certain parts aren't known. + # TODO: In this case, the reified Theano object is a sort of + # "proxy" object; we should use this approach for dtype, as well. + # TODO: We should also put this kind of logic in the appropriate places + # (e.g. `MetaVariable.reify`), when possible. + if TheanoMetaSymbol.is_meta(name): + # This should also invalidate `res_var.obj`. + res_var.name = name + # Allow the base object to be unified, so that reification + # can recover the underlying object--instead of recreating + # it and sacrificing equality. + res_var.obj = var() + + elif tt_out.name != name: + tt_out.name = name + res_var.name = name + + else: + # XXX: It's not always clear how `Op.make_node` arguments map to + # `Apply` node inputs, which is one of the big problem with + # Theano's design. (More generally, it's that `Op`s don't provide + # a spec for `Apply` node inputs and outputs at all.) + + # Also, `Apply` inputs can't be `None` (they could be + # `tt.none_type_t()`, though). + res_apply = TheanoMetaApply( + self, tuple(filter(lambda x: x is not None, op_arg_bind.args)) + ) + + # TODO: Elemwise has an `output_types` method that can be + # used to infer the output type of this variable. + ttype = ttype or var() + + # Use the given index or the base `Op`'s `default_output`; + # otherwise, create a logic variable place-holder. + index = index if index is not None else getattr(self.obj, "default_output", None) + index = index if index is not None else var() + + # XXX: We don't have a higher-order meta object model, so being + # wrong about the exact type of output variable will cause + # problems. + out_meta_type = self.out_meta_type(op_args) + res_var = out_meta_type(ttype, res_apply, index, name) + res_var.obj = var() + + return res_var + + +class TheanoMetaElemwise(TheanoMetaOp): + base = tt.Elemwise + + def __call__(self, *args, ttype=None, index=None, **kwargs): + obj_nout = getattr(self.obj, "nfunc_spec", None) + obj_nout = obj_nout[-1] if obj_nout is not None else None + if obj_nout == 1 and index is None: + index = 0 + return super().__call__(*args, ttype=ttype, index=index, **kwargs) + + +class TheanoMetaDimShuffle(TheanoMetaOp): + base = tt.DimShuffle + __slots__ = ["input_broadcastable", "new_order", "inplace"] + + def __init__(self, input_broadcastable, new_order, inplace=True, obj=None): + super().__init__(obj=obj) + self.input_broadcastable = input_broadcastable + self.new_order = new_order + self.inplace = inplace + + +class TheanoMetaRandomVariable(TheanoMetaOp): + base = RandomVariable + + def __init__(self, obj=None): + super().__init__(obj=obj) + # The `name` keyword parameter isn't an `Apply` node input, so we need + # to remove it from the automatically generated signature. + self.op_sig = self.op_sig.replace(parameters=list(self.op_sig.parameters.values())[0:4]) + + +class TheanoMetaApply(TheanoMetaSymbol): + base = tt.Apply + __slots__ = ["op", "inputs"] + + def __init__(self, op, inputs, outputs=None, obj=None): + super().__init__(obj=obj) + self.op = metatize(op) + self.inputs = tuple(metatize(i) for i in inputs) + self.outputs = outputs + + def reify(self): + if self.obj and not isinstance(self.obj, Var): + return self.obj + else: + tt_op = self.op.reify() + if not self.is_meta(tt_op): + reified_rands, any_unreified = _meta_reify_iter(self.inputs) + if not any_unreified: + tt_var = tt_op(*reified_rands) + self.obj = tt_var.owner + return tt_var.owner + return self + + @property + def nin(self): + return len(self.inputs) + + @property + def nout(self): + if self.outputs is not None: + return len(self.outputs) + elif self.obj: + return len(self.obj.outputs) + # TODO: Would be cool if we could return + # a logic variable representing this. + + +class TheanoMetaVariable(MetaVariable, TheanoMetaSymbol): + base = theano.Variable + __slots__ = ["type", "owner", "index", "name"] + + def __init__(self, type, owner, index, name, obj=None): + super().__init__(obj=obj) + self.type = metatize(type) + self.owner = metatize(owner) + self.index = index + self.name = name + + def reify(self): + if self.obj and not isinstance(self.obj, Var): + return self.obj + + if not self.owner: + return super().reify() + + # Having an `owner` causes issues (e.g. being consistent about + # other, unrelated outputs of an `Apply` node), and, in this case, + # the `Apply` node that owns this variable needs to construct it. + reified_rands, any_unreified = _meta_reify_iter(self.rands()) + tt_apply = self.owner.obj + + if tt_apply and not isvar(tt_apply): + # If the owning `Apply` reified, then one of its `outputs` + # corresponds to this variable. Our `self.index` value should + # tell us which, but, when that's not available, we can + # sometimes infer it. + if tt_apply.nout == 1: + tt_index = 0 + # Make sure we didn't have a mismatched non-meta index value. + assert isvar(self.index) or self.index is None or self.index == 0 + # Set/replace `None` or meta value + self.index = 0 + tt_var = tt_apply.outputs[tt_index] + elif not self.is_meta(self.index): + tt_var = tt_apply.outputs[self.index] + elif self.index is None: + tt_var = tt_apply.default_output() + self.index = tt_apply.outputs.index(tt_var) + else: + return self + # If our name value is not set/concrete, then use the reified + # value's. Otherwise, use ours. + if isvar(self.name) or self.name is None: + self.name = tt_var.name + else: + tt_var.name = self.name + assert tt_var is not None + self.obj = tt_var + return tt_var + + return super().reify() + + +class TheanoMetaTensorVariable(TheanoMetaVariable): + # TODO: Could extend `theano.tensor.var._tensor_py_operators`, too. + base = tt.TensorVariable + + @property + def ndim(self): + if isinstance(self.type, TheanoMetaTensorType) and isinstance( + self.type.broadastable, (list, tuple) + ): + return len(self.type.broadcastable) + # TODO: Would be cool if we could return + # a logic variable representing this. + + +class TheanoMetaConstant(TheanoMetaVariable): + base = theano.Constant + __slots__ = ["type", "data"] + + def __init__(self, type, data, name=None, obj=None): + super().__init__(type, None, None, name, obj=obj) + self.data = data + + +class TheanoMetaTensorConstant(TheanoMetaConstant): + # TODO: Could extend `theano.tensor.var._tensor_py_operators`, too. + base = tt.TensorConstant + __slots__ = ["type", "data", "name"] + + def __init__(self, type, data, name=None, obj=None): + super().__init__(type, data, name, obj=obj) + + +class TheanoMetaSharedVariable(TheanoMetaVariable): + base = tt.sharedvar.SharedVariable + __slots__ = ["name", "type", "data", "strict"] + + @classmethod + def _metatize(cls, obj): + res = TheanoMetaSharedVariable( + obj.name, obj.type, obj.container.data, obj.container.strict, obj=obj + ) + return res + + def __init__(self, name, type, data, strict, obj=None): + super().__init__(type, None, None, name, obj=obj) + self.data = data + self.strict = strict + + +class TheanoMetaTensorSharedVariable(TheanoMetaSharedVariable): + # TODO: Could extend `theano.tensor.var._tensor_py_operators`, too. + base = tt.sharedvar.TensorSharedVariable + + +class TheanoMetaScalarSharedVariable(TheanoMetaSharedVariable): + base = tt.sharedvar.ScalarSharedVariable + + +class TheanoMetaAccessor(object): + """Create an object that can be used to implicitly convert Theano functions and object into meta objects. + + Use it like a namespace/module/package object, e.g. + + >>> mt = TheanoMetaAccessor() + >>> mt.vector('a') + MetaTensorVariable(MetaTensorType('float64', (False,), None, + obj=TensorType(float64, vector)), None, None, 'a', obj=a) + + Call it as a function to perform direct conversion to a meta + object, e.g. + + >>> mt(tt.vector('a')) + MetaTensorVariable(MetaTensorType('float64', (False,), None, + obj=TensorType(float64, vector)), None, None, 'a', obj=a) + + """ + + namespaces = [tt] + + def __init__(self, namespace=None): + if namespace is None: + from symbolic_pymc.theano import ( # pylint: disable=import-self + meta, + ops, + random_variables, + ) + + self.namespaces += [meta, ops, random_variables] + else: + self.namespaces = [namespace] + + def __call__(self, x): + return metatize(x) + + def __getattr__(self, obj): + + ns_obj = next((getattr(ns, obj) for ns in self.namespaces if hasattr(ns, obj)), None) + + if ns_obj is None: + # Try caller's namespace + frame = inspect.currentframe() + f_back = frame.f_back + if f_back: + ns_obj = f_back.f_locals.get(obj, None) + if ns_obj is None: + ns_obj = f_back.f_globals.get(obj) + + if isinstance(ns_obj, (types.FunctionType, partial)): + # It's a function, so let's provide a wrapper that converts + # to-and-from theano and meta objects. + @wraps(ns_obj) + def meta_obj(*args, **kwargs): + args = [o.reify() if hasattr(o, "reify") else o for o in args] + res = ns_obj(*args, **kwargs) + return metatize(res) + + elif isinstance(ns_obj, types.ModuleType): + # It's a sub-module, so let's create another + # `TheanoMetaAccessor` and check within there. + meta_obj = TheanoMetaAccessor(namespace=ns_obj) + else: + # Hopefully, it's convertible to a meta object... + meta_obj = metatize(ns_obj) + + if isinstance(meta_obj, (TheanoMetaSymbol, MetaSymbolType, types.FunctionType)): + setattr(self, obj, meta_obj) + return getattr(self, obj) + elif isinstance(meta_obj, TheanoMetaAccessor): + setattr(self, obj, meta_obj) + return meta_obj + else: + raise AttributeError(f"Meta object for {obj} not found.") + + +mt = TheanoMetaAccessor() + +mt.dot = metatize(tt.basic._dot) + + +# +# The wrapped Theano functions will only work when the meta objects +# are fully reifiable (i.e. can be turned to Theano objects), but it's +# fairly straight-forward to adjust many of those functions so that they +# work with meta objects. +# TODO: Would be nice if we could trick Theano into using meta objects, or +# a robust use of "proxy" Theano objects +# + + +def mt_zeros(shape, dtype=None): + if not isinstance(shape, (list, tuple, TheanoMetaTensorVariable, tt.TensorVariable)): + shape = [shape] + if dtype is None: + dtype = tt.config.floatX + return mt.alloc(np.array(0, dtype=dtype), *shape) + + +mt.zeros = mt_zeros + + +def mt_diag(v, k=0): + if v.ndim == 1: + return mt.AllocDiag(k)(v) + elif v.ndim is not None and v.ndim >= 2: + return mt.diagonal(v, offset=k) + else: + raise ValueError("Input must has v.ndim >= 1.") + + +fact(commutative, mt.add) +fact(commutative, mt.mul) +fact(associative, mt.add) +fact(associative, mt.mul) diff --git a/symbolic_pymc/rv.py b/symbolic_pymc/theano/ops.py similarity index 100% rename from symbolic_pymc/rv.py rename to symbolic_pymc/theano/ops.py diff --git a/symbolic_pymc/opt.py b/symbolic_pymc/theano/opt.py similarity index 100% rename from symbolic_pymc/opt.py rename to symbolic_pymc/theano/opt.py diff --git a/symbolic_pymc/printing.py b/symbolic_pymc/theano/printing.py similarity index 99% rename from symbolic_pymc/printing.py rename to symbolic_pymc/theano/printing.py index b28b9e0..a4f9505 100644 --- a/symbolic_pymc/printing.py +++ b/symbolic_pymc/theano/printing.py @@ -13,9 +13,9 @@ from sympy import Array as SympyArray from sympy.printing import latex as sympy_latex -from . import Observed, NormalRV from .opt import FunctionGraph -from .rv import RandomVariable +from .ops import RandomVariable +from .random_variables import Observed, NormalRV class RandomVariablePrinter(object): diff --git a/symbolic_pymc/pymc3.py b/symbolic_pymc/theano/pymc3.py similarity index 97% rename from symbolic_pymc/pymc3.py rename to symbolic_pymc/theano/pymc3.py index 755748c..54d42f6 100644 --- a/symbolic_pymc/pymc3.py +++ b/symbolic_pymc/theano/pymc3.py @@ -17,7 +17,7 @@ from theano.gof.graph import Apply, inputs as tt_inputs -from . import ( +from .random_variables import ( observed, UniformRV, UniformRVType, @@ -39,7 +39,7 @@ HalfCauchyRVType, ) from .opt import FunctionGraph -from .rv import RandomVariable +from .ops import RandomVariable from .utils import replace_input_nodes, get_rv_observation logger = logging.getLogger("symbolic_pymc") @@ -96,7 +96,8 @@ def _convert_rv_to_dist(op, rv): @dispatch(pm.HalfNormal, object) def convert_dist_to_rv(dist, rng): size = dist.shape.astype(int)[HalfNormalRV.ndim_supp :] - res = HalfNormalRV(0.0, dist.sd, size=size, rng=rng) + res = HalfNormalRV(np.array(0.0, dtype=dist.dtype), + dist.sd, size=size, rng=rng) return res @@ -175,7 +176,8 @@ def _convert_rv_to_dist(op, rv): @dispatch(pm.HalfCauchy, object) def convert_dist_to_rv(dist, rng): size = dist.shape.astype(int)[HalfCauchyRV.ndim_supp :] - res = HalfCauchyRV(0.0, dist.beta, size=size, rng=rng) + res = HalfCauchyRV(np.array(0.0, dtype=dist.dtype), + dist.beta, size=size, rng=rng) return res diff --git a/symbolic_pymc/theano/random_variables.py b/symbolic_pymc/theano/random_variables.py new file mode 100644 index 0000000..0b80036 --- /dev/null +++ b/symbolic_pymc/theano/random_variables.py @@ -0,0 +1,280 @@ +import theano +import scipy +import theano.tensor as tt + +from functools import partial + +from .ops import RandomVariable, param_supp_shape_fn + + +# Continuous Numpy-generated variates +class UniformRVType(RandomVariable): + print_name = ("U", "\\operatorname{U}") + + def __init__(self): + super().__init__("uniform", theano.config.floatX, 0, [0, 0], "uniform", inplace=True) + + def make_node(self, lower, upper, size=None, rng=None, name=None): + return super().make_node(lower, upper, size=size, rng=rng, name=name) + + +UniformRV = UniformRVType() + + +class NormalRVType(RandomVariable): + print_name = ("N", "\\operatorname{N}") + + def __init__(self): + super().__init__("normal", theano.config.floatX, 0, [0, 0], "normal", inplace=True) + + def make_node(self, mu, sigma, size=None, rng=None, name=None): + return super().make_node(mu, sigma, size=size, rng=rng, name=name) + + +NormalRV = NormalRVType() + + +class HalfNormalRVType(RandomVariable): + print_name = ("N**+", "\\operatorname{N^{+}}") + + def __init__(self): + super().__init__( + "halfnormal", + theano.config.floatX, + 0, + [0, 0], + lambda rng, *args: scipy.stats.halfnorm.rvs(*args, random_state=rng), + inplace=True, + ) + + def make_node(self, mu=0.0, sigma=1.0, size=None, rng=None, name=None): + return super().make_node(mu, sigma, size=size, rng=rng, name=name) + + +HalfNormalRV = HalfNormalRVType() + + +class GammaRVType(RandomVariable): + print_name = ("Gamma", "\\operatorname{Gamma}") + + def __init__(self): + super().__init__("gamma", theano.config.floatX, 0, [0, 0], "gamma", inplace=True) + + def make_node(self, shape, scale, size=None, rng=None, name=None): + return super().make_node(shape, scale, size=size, rng=rng, name=name) + + +GammaRV = GammaRVType() + + +class ExponentialRVType(RandomVariable): + print_name = ("Exp", "\\operatorname{Exp}") + + def __init__(self): + super().__init__("exponential", theano.config.floatX, 0, [0], "exponential", inplace=True) + + def make_node(self, scale, size=None, rng=None, name=None): + return super().make_node(scale, size=size, rng=rng, name=name) + + +ExponentialRV = ExponentialRVType() + + +# One with multivariate support +class MvNormalRVType(RandomVariable): + print_name = ("N", "\\operatorname{N}") + + def __init__(self): + super().__init__( + "multivariate_normal", + theano.config.floatX, + 1, + [1, 2], + "multivariate_normal", + inplace=True, + ) + + def make_node(self, mean, cov, size=None, rng=None, name=None): + return super().make_node(mean, cov, size=size, rng=rng, name=name) + + +MvNormalRV = MvNormalRVType() + + +class DirichletRVType(RandomVariable): + print_name = ("Dir", "\\operatorname{Dir}") + + def __init__(self): + super().__init__("dirichlet", theano.config.floatX, 1, [1], "dirichlet", inplace=True) + + def make_node(self, alpha, size=None, rng=None, name=None): + return super().make_node(alpha, size=size, rng=rng, name=name) + + +DirichletRV = DirichletRVType() + + +# A discrete Numpy-generated variate +class PoissonRVType(RandomVariable): + print_name = ("Pois", "\\operatorname{Pois}") + + def __init__(self): + super().__init__("poisson", "int64", 0, [0], "poisson", inplace=True) + + def make_node(self, rate, size=None, rng=None, name=None): + return super().make_node(rate, size=size, rng=rng, name=name) + + +PoissonRV = PoissonRVType() + + +# A SciPy-generated variate +class CauchyRVType(RandomVariable): + print_name = ("C", "\\operatorname{C}") + + def __init__(self): + super().__init__( + "cauchy", + theano.config.floatX, + 0, + [0, 0], + lambda rng, *args: scipy.stats.cauchy.rvs(*args, random_state=rng), + inplace=True, + ) + + def make_node(self, loc, scale, size=None, rng=None, name=None): + return super().make_node(loc, scale, size=size, rng=rng, name=name) + + +CauchyRV = CauchyRVType() + + +class HalfCauchyRVType(RandomVariable): + print_name = ("C**+", "\\operatorname{C^{+}}") + + def __init__(self): + super().__init__( + "halfcauchy", + theano.config.floatX, + 0, + [0, 0], + lambda rng, *args: scipy.stats.halfcauchy.rvs(*args, random_state=rng), + inplace=True, + ) + + def make_node(self, loc=0.0, scale=1.0, size=None, rng=None, name=None): + return super().make_node(loc, scale, size=size, rng=rng, name=name) + + +HalfCauchyRV = HalfCauchyRVType() + + +class InvGammaRVType(RandomVariable): + print_name = ("InvGamma", "\\operatorname{Gamma^{-1}}") + + def __init__(self): + super().__init__( + "invgamma", + theano.config.floatX, + 0, + [0, 0], + lambda rng, *args: scipy.stats.invgamma.rvs(*args, random_state=rng), + inplace=True, + ) + + def make_node(self, loc, scale, size=None, rng=None, name=None): + return super().make_node(loc, scale, size=size, rng=rng, name=name) + + +InvGammaRV = InvGammaRVType() + + +class TruncExponentialRVType(RandomVariable): + print_name = ("TruncExp", "\\operatorname{Exp}") + + def __init__(self): + super().__init__( + "truncexpon", + theano.config.floatX, + 0, + [0, 0, 0], + lambda rng, *args: scipy.stats.truncexpon.rvs(*args, random_state=rng), + inplace=True, + ) + + def make_node(self, b, loc, scale, size=None, rng=None, name=None): + return super().make_node(b, loc, scale, size=size, rng=rng, name=name) + + +TruncExponentialRV = TruncExponentialRVType() + + +# Support shape is determined by the first dimension in the *second* parameter +# (i.e. the probabilities vector) +class MultinomialRVType(RandomVariable): + print_name = ("MN", "\\operatorname{MN}") + + def __init__(self): + super().__init__( + "multinomial", + "int64", + 1, + [0, 1], + "multinomial", + supp_shape_fn=partial(param_supp_shape_fn, rep_param_idx=1), + inplace=True, + ) + + def make_node(self, n, pvals, size=None, rng=None, name=None): + return super().make_node(n, pvals, size=size, rng=rng, name=name) + + +MultinomialRV = MultinomialRVType() + + +class Observed(tt.Op): + """An `Op` that represents an observed random variable. + + This `Op` establishes an observation relationship between a random + variable and a specific value. + """ + + default_output = 0 + + def __init__(self): + self.view_map = {0: [0]} + + def make_node(self, val, rv=None): + """Make an `Observed` random variable. + + Parameters + ---------- + val: Variable + The observed value. + rv: RandomVariable + The distribution from which `val` is assumed to be a sample value. + """ + val = tt.as_tensor_variable(val) + if rv: + if rv.owner and not isinstance(rv.owner.op, RandomVariable): + raise ValueError(f"`rv` must be a RandomVariable type: {rv}") + + if rv.type.convert_variable(val) is None: + raise ValueError( + ("`rv` and `val` do not have compatible types:" f" rv={rv}, val={val}") + ) + else: + rv = tt.NoneConst.clone() + + inputs = [val, rv] + + return tt.Apply(self, inputs, [val.type()]) + + def perform(self, node, inputs, out): + out[0][0] = inputs[0] + + def grad(self, inputs, outputs): + return outputs + + +observed = Observed() diff --git a/symbolic_pymc/theano/unify.py b/symbolic_pymc/theano/unify.py new file mode 100644 index 0000000..38a8b90 --- /dev/null +++ b/symbolic_pymc/theano/unify.py @@ -0,0 +1,47 @@ +import theano.tensor as tt + +from multipledispatch import dispatch + +from kanren.term import term, operator, arguments + +from unification.core import _reify, _unify, reify + +from ..meta import metatize +from ..unify import ExpressionTuple, unify_MetaSymbol +from .meta import mt, TheanoMetaSymbol + + +tt_class_abstractions = tuple(c.base for c in TheanoMetaSymbol.__subclasses__()) + + +_unify.add( + (TheanoMetaSymbol, tt_class_abstractions, dict), + lambda u, v, s: unify_MetaSymbol(u, metatize(v), s), +) +_unify.add( + (tt_class_abstractions, TheanoMetaSymbol, dict), + lambda u, v, s: unify_MetaSymbol(metatize(u), v, s), +) +_unify.add( + (tt_class_abstractions, tt_class_abstractions, dict), + lambda u, v, s: unify_MetaSymbol(metatize(u), metatize(v), s), +) + + +def _reify_TheanoClasses(o, s): + meta_obj = metatize(o) + return reify(meta_obj, s) + + +_reify.add((tt_class_abstractions, dict), _reify_TheanoClasses) + +operator.add((tt.Variable,), lambda x: operator(metatize(x))) + +arguments.add((tt.Variable,), lambda x: arguments(metatize(x))) + +term.add((tt.Op, ExpressionTuple), lambda op, args: term(metatize(op), args)) + + +@dispatch(tt_class_abstractions) +def tuple_expression(x): + return tuple_expression(mt(x)) diff --git a/symbolic_pymc/theano/utils.py b/symbolic_pymc/theano/utils.py new file mode 100644 index 0000000..0dc9b4b --- /dev/null +++ b/symbolic_pymc/theano/utils.py @@ -0,0 +1,147 @@ +import theano.tensor as tt + +from theano.gof import FunctionGraph as tt_FunctionGraph, Query +from theano.gof.graph import inputs as tt_inputs, clone_get_equiv, io_toposort +from theano.compile import optdb + +from .meta import mt +from .opt import FunctionGraph +from .ops import RandomVariable +from .random_variables import Observed + +from unification.utils import transitive_get as walk + + +canonicalize_opt = optdb.query(Query(include=["canonicalize"])) + + +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 + + +def graph_equal(x, y): + """Compare elements in a Theano graph using their object properties and not just identity.""" + try: + if isinstance(x, (list, tuple)) and isinstance(y, (list, tuple)): + return len(x) == len(y) and all(mt(xx) == mt(yy) for xx, yy in zip(x, y)) + return mt(x) == mt(y) + except ValueError: + return False + + +def optimize_graph(x, optimization, return_graph=None, in_place=False): + """Easily optimize Theano graphs. + + Apply an optimization to either the graph formed by a Theano variable or an + existing graph and return the resulting optimized graph. + + When given an existing `FunctionGraph`, the optimization is + performed without side-effects (i.e. won't change the given graph). + + """ + if not isinstance(x, tt_FunctionGraph): + inputs = tt_inputs([x]) + outputs = [x] + model_memo = clone_get_equiv(inputs, outputs, copy_orphans=False) + cloned_inputs = [model_memo[i] for i in inputs if not isinstance(i, tt.Constant)] + cloned_outputs = [model_memo[i] for i in outputs] + + x_graph = FunctionGraph(cloned_inputs, cloned_outputs, clone=False) + x_graph.memo = model_memo + + if return_graph is None: + return_graph = False + else: + x_graph = x + + if return_graph is None: + return_graph = True + + x_graph_opt = x_graph if in_place else x_graph.clone() + _ = optimization.optimize(x_graph_opt) + + if return_graph: + res = x_graph_opt + else: + res = x_graph_opt.outputs + if len(res) == 1: + res, = res + return res + + +def canonicalize(x, **kwargs): + """Canonicalize a Theano variable and/or graph.""" + return optimize_graph(x, canonicalize_opt, **kwargs) + + +def get_rv_observation(node): + """Return a `RandomVariable` node's corresponding `Observed` node, or `None`.""" + if not getattr(node, "fgraph", None): + raise ValueError("Node does not belong to a `FunctionGraph`") + + if isinstance(node.op, RandomVariable): + fgraph = node.fgraph + for o, i in node.default_output().clients: + if o == "output": + o = fgraph.outputs[i].owner + + if isinstance(o.op, Observed): + return o + return None diff --git a/symbolic_pymc/unify.py b/symbolic_pymc/unify.py index 0bfb7f1..978d11f 100644 --- a/symbolic_pymc/unify.py +++ b/symbolic_pymc/unify.py @@ -1,26 +1,18 @@ import reprlib -import theano.tensor as tt - from functools import wraps from multipledispatch import dispatch from kanren import isvar from kanren.term import term, operator, arguments -from kanren.facts import fact -from kanren.assoccomm import commutative, associative - -# from kanren.goals import LCons, _lcons_unify from unification.more import unify from unification.core import reify, _unify, _reify, Var -from .meta import MetaSymbol, MetaVariable, mt, metatize +from .meta import MetaSymbol, MetaVariable from .utils import _check_eq -tt_class_abstractions = tuple(c.base for c in MetaSymbol.__subclasses__()) - etuple_repr = reprlib.Repr() etuple_repr.maxstring = 100 etuple_repr.maxother = 100 @@ -171,16 +163,6 @@ def unify_MetaSymbol(u, v, s): _unify.add((MetaSymbol, MetaSymbol, dict), unify_MetaSymbol) -_unify.add( - (MetaSymbol, tt_class_abstractions, dict), lambda u, v, s: unify_MetaSymbol(u, metatize(v), s) -) -_unify.add( - (tt_class_abstractions, MetaSymbol, dict), lambda u, v, s: unify_MetaSymbol(metatize(u), v, s) -) -_unify.add( - (tt_class_abstractions, tt_class_abstractions, dict), - lambda u, v, s: unify_MetaSymbol(metatize(u), metatize(v), s), -) def _reify_MetaSymbol(o, s): @@ -202,14 +184,6 @@ def _reify_MetaSymbol(o, s): _reify.add((MetaSymbol, dict), _reify_MetaSymbol) -def _reify_TheanoClasses(o, s): - meta_obj = metatize(o) - return reify(meta_obj, s) - - -_reify.add((tt_class_abstractions, dict), _reify_TheanoClasses) - - _isvar = isvar.dispatch(object) isvar.add((MetaSymbol,), lambda x: _isvar(x) or (not isinstance(x.obj, Var) and isvar(x.obj))) @@ -254,7 +228,6 @@ def operator_MetaVariable(x): operator.add((MetaSymbol,), operator_MetaSymbol) operator.add((MetaVariable,), operator_MetaVariable) -operator.add((tt.Variable,), lambda x: operator(metatize(x))) def arguments_MetaSymbol(x): @@ -289,7 +262,6 @@ def arguments_MetaVariable(x): arguments.add((MetaSymbol,), arguments_MetaSymbol) arguments.add((MetaVariable,), arguments_MetaVariable) -arguments.add((tt.Variable,), lambda x: arguments(metatize(x))) def _term_ExpressionTuple(rand, rators): @@ -298,7 +270,6 @@ def _term_ExpressionTuple(rand, rators): term.add((object, ExpressionTuple), _term_ExpressionTuple) -term.add((tt.Op, ExpressionTuple), lambda op, args: term(metatize(op), args)) @dispatch(object) @@ -337,11 +308,6 @@ def tuple_expression(x): return res -@dispatch(tt_class_abstractions) -def tuple_expression(x): - return tuple_expression(mt(x)) - - @_reify.register(ExpressionTuple, dict) def _reify_ExpressionTuple(t, s): """When `kanren` reifies `etuple`s, we don't want them to turn into regular `tuple`s. @@ -370,9 +336,4 @@ def _reify_ExpressionTuple(t, s): return res -fact(commutative, mt.add) -fact(commutative, mt.mul) -fact(associative, mt.add) -fact(associative, mt.mul) - __all__ = ["debug_unify", "etuple", "tuple_expression"] diff --git a/symbolic_pymc/utils.py b/symbolic_pymc/utils.py index 5f19730..dbfcf81 100644 --- a/symbolic_pymc/utils.py +++ b/symbolic_pymc/utils.py @@ -1,21 +1,7 @@ -import theano -import theano.tensor as tt - import numpy as np import symbolic_pymc as sp -from collections import OrderedDict - -from unification.utils import transitive_get as walk - -from theano.gof import FunctionGraph as tt_FunctionGraph, Query -from theano.gof.graph import inputs as tt_inputs, clone_get_equiv, io_toposort -from theano.compile import optdb - - -canonicalize_opt = optdb.query(Query(include=["canonicalize"])) - def _check_eq(a, b): if isinstance(a, np.ndarray) or isinstance(b, np.ndarray): @@ -24,83 +10,6 @@ def _check_eq(a, b): return a == b -def get_rv_observation(node): - """Return a `RandomVariable` node's corresponding `Observed` node, or `None`.""" - if not getattr(node, "fgraph", None): - raise ValueError("Node does not belong to a `FunctionGraph`") - - if isinstance(node.op, sp.rv.RandomVariable): - fgraph = node.fgraph - for o, i in node.default_output().clients: - if o == "output": - o = fgraph.outputs[i].owner - - if isinstance(o.op, sp.Observed): - return o - return None - - -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 - - def meta_parts_unequal(x, y, pdb=False): """Traverse meta objects and return the first pair of elements that are not equal.""" res = None @@ -132,80 +41,3 @@ def meta_parts_unequal(x, y, pdb=False): pdb.set_trace() return res - - -def expand_meta(x, tt_print=tt.pprint): - """Produce a dictionary representation of a meta object.""" - if isinstance(x, sp.meta.MetaSymbol): - return OrderedDict( - [ - ("rator", x.base), - ("rands", tuple(expand_meta(p) for p in x.rands())), - ("obj", expand_meta(getattr(x, "obj", None))), - ] - ) - elif tt_print and isinstance(x, theano.gof.op.Op): - return x.name - elif tt_print and isinstance(x, theano.gof.graph.Variable): - return tt_print(x) - else: - return x - - -def graph_equal(x, y): - """Compare elements in a Theano graph using their object properties and not just identity.""" - try: - if isinstance(x, (list, tuple)) and isinstance(y, (list, tuple)): - return len(x) == len(y) and all(sp.mt(xx) == sp.mt(yy) for xx, yy in zip(x, y)) - return sp.mt(x) == sp.mt(y) - except ValueError: - return False - - -def mt_type_params(x): - return {"ttype": x.type, "index": x.index, "name": x.name} - - -def optimize_graph(x, optimization, return_graph=None, in_place=False): - """Easily optimize Theano graphs. - - Apply an optimization to either the graph formed by a Theano variable or an - existing graph and return the resulting optimized graph. - - When given an existing `FunctionGraph`, the optimization is - performed without side-effects (i.e. won't change the given graph). - - """ - if not isinstance(x, tt_FunctionGraph): - inputs = tt_inputs([x]) - outputs = [x] - model_memo = clone_get_equiv(inputs, outputs, copy_orphans=False) - cloned_inputs = [model_memo[i] for i in inputs if not isinstance(i, tt.Constant)] - cloned_outputs = [model_memo[i] for i in outputs] - - x_graph = sp.opt.FunctionGraph(cloned_inputs, cloned_outputs, clone=False) - x_graph.memo = model_memo - - if return_graph is None: - return_graph = False - else: - x_graph = x - - if return_graph is None: - return_graph = True - - x_graph_opt = x_graph if in_place else x_graph.clone() - _ = optimization.optimize(x_graph_opt) - - if return_graph: - res = x_graph_opt - else: - res = x_graph_opt.outputs - if len(res) == 1: - res, = res - return res - - -def canonicalize(x, **kwargs): - """Canonicalize a Theano variable and/or graph.""" - return optimize_graph(x, canonicalize_opt, **kwargs) diff --git a/tests/__init__.py b/tests/__init__.py index 70ecd26..e69de29 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -1,13 +0,0 @@ -import warnings - -with warnings.catch_warnings(): - warnings.filterwarnings("ignore", category=DeprecationWarning) - import theano - -import pymc3 as pm - - -theano.config.compute_test_value = 'ignore' -theano.config.on_opt_error = 'raise' -theano.config.mode = 'FAST_COMPILE' -theano.config.cxx = '' diff --git a/tests/test_relations.py b/tests/test_relations.py deleted file mode 100644 index 7aa5ace..0000000 --- a/tests/test_relations.py +++ /dev/null @@ -1,82 +0,0 @@ -import numpy as np -import theano -import theano.tensor as tt - -from theano.gof.opt import EquilibriumOptimizer - -from symbolic_pymc import (observed, NormalRV, HalfCauchyRV) -from symbolic_pymc.opt import KanrenRelationSub, FunctionGraph -from symbolic_pymc.utils import (optimize_graph, canonicalize, - get_rv_observation) -from symbolic_pymc.relations.distributions import scale_loc_transform - - -def test_pymc_normals(): - tt.config.compute_test_value = 'ignore' - - rand_state = theano.shared(np.random.RandomState()) - mu_a = NormalRV(0., 100**2, name='mu_a', rng=rand_state) - sigma_a = HalfCauchyRV(5, name='sigma_a', rng=rand_state) - mu_b = NormalRV(0., 100**2, name='mu_b', rng=rand_state) - sigma_b = HalfCauchyRV(5, name='sigma_b', rng=rand_state) - county_idx = np.r_[1, 1, 2, 3] - # We want the following for a, b: - # N(m, S) -> m + N(0, 1) * S - a = NormalRV(mu_a, sigma_a, size=(len(county_idx),), name='a', rng=rand_state) - b = NormalRV(mu_b, sigma_b, size=(len(county_idx),), name='b', rng=rand_state) - radon_est = a[county_idx] + b[county_idx] * 7 - eps = HalfCauchyRV(5, name='eps', rng=rand_state) - radon_like = NormalRV(radon_est, eps, name='radon_like', rng=rand_state) - radon_like_rv = observed(tt.as_tensor_variable(np.r_[1., 2., 3., 4.]), radon_like) - - inputs = [mu_a, mu_b, eps, rand_state] - fgraph = FunctionGraph( - inputs, - [radon_like_rv], - clone=True) - fgraph = canonicalize(fgraph, return_graph=True, in_place=False) - - posterior_opt = EquilibriumOptimizer( - [KanrenRelationSub(scale_loc_transform, - node_filter=get_rv_observation)], - max_use_ratio=10) - - fgraph_opt = optimize_graph(fgraph, posterior_opt, return_graph=True) - fgraph_opt = canonicalize(fgraph_opt, return_graph=True, in_place=False) - - radon_like_rv_opt = fgraph_opt.outputs[0] - radon_like_opt = radon_like_rv_opt.owner.inputs[1] - radon_est_opt = radon_like_opt.owner.inputs[0] - - # These should now be `tt.add(mu_*, ...)` outputs. - a_opt = radon_est_opt.owner.inputs[0].owner.inputs[0] - b_opt = radon_est_opt.owner.inputs[1].owner.inputs[1].owner.inputs[0] - - # Make sure NormalRV gets replaced with an addition - assert a_opt.owner.op == tt.add - assert b_opt.owner.op == tt.add - - # Make sure the first term in the addition is the old NormalRV mean - mu_a_opt = a_opt.owner.inputs[0].owner.inputs[0] - assert 'mu_a' == mu_a_opt.name == mu_a.name - mu_b_opt = b_opt.owner.inputs[0].owner.inputs[0] - assert 'mu_b' == mu_b_opt.name == mu_b.name - - # Make sure the second term in the addition is the standard NormalRV times - # the old std. dev. - assert a_opt.owner.inputs[1].owner.op == tt.mul - assert b_opt.owner.inputs[1].owner.op == tt.mul - - sigma_a_opt = a_opt.owner.inputs[1].owner.inputs[0].owner.inputs[0] - assert sigma_a_opt.owner.op == sigma_a.owner.op - sigma_b_opt = b_opt.owner.inputs[1].owner.inputs[0].owner.inputs[0] - assert sigma_b_opt.owner.op == sigma_b.owner.op - - a_std_norm_opt = a_opt.owner.inputs[1].owner.inputs[1] - assert a_std_norm_opt.owner.op == NormalRV - assert a_std_norm_opt.owner.inputs[0].data == 0.0 - assert a_std_norm_opt.owner.inputs[1].data == 1.0 - b_std_norm_opt = b_opt.owner.inputs[1].owner.inputs[1] - assert b_std_norm_opt.owner.op == NormalRV - assert b_std_norm_opt.owner.inputs[0].data == 0.0 - assert b_std_norm_opt.owner.inputs[1].data == 1.0 diff --git a/tests/theano/__init__.py b/tests/theano/__init__.py new file mode 100644 index 0000000..70ecd26 --- /dev/null +++ b/tests/theano/__init__.py @@ -0,0 +1,13 @@ +import warnings + +with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=DeprecationWarning) + import theano + +import pymc3 as pm + + +theano.config.compute_test_value = 'ignore' +theano.config.on_opt_error = 'raise' +theano.config.mode = 'FAST_COMPILE' +theano.config.cxx = '' diff --git a/tests/test_conjugates.py b/tests/theano/test_conjugates.py similarity index 91% rename from tests/test_conjugates.py rename to tests/theano/test_conjugates.py index 3474d8f..d801baa 100644 --- a/tests/test_conjugates.py +++ b/tests/theano/test_conjugates.py @@ -4,9 +4,9 @@ from theano.gof.opt import EquilibriumOptimizer from theano.gof.graph import inputs as tt_inputs -from symbolic_pymc import MvNormalRV, observed -from symbolic_pymc.opt import KanrenRelationSub, FunctionGraph -from symbolic_pymc.utils import optimize_graph +from symbolic_pymc.theano.random_variables import MvNormalRV, observed +from symbolic_pymc.theano.opt import KanrenRelationSub, FunctionGraph +from symbolic_pymc.theano.utils import optimize_graph from symbolic_pymc.relations.conjugates import conjugate_posteriors diff --git a/tests/test_kanren.py b/tests/theano/test_kanren.py similarity index 96% rename from tests/test_kanren.py rename to tests/theano/test_kanren.py index 367bab0..e19656a 100644 --- a/tests/test_kanren.py +++ b/tests/theano/test_kanren.py @@ -5,10 +5,10 @@ from kanren.assoccomm import eq_assoc, eq_comm from unification import var -from symbolic_pymc import MvNormalRV -from symbolic_pymc.meta import mt +from symbolic_pymc.theano.random_variables import MvNormalRV +from symbolic_pymc.theano.meta import mt from symbolic_pymc.unify import etuple -from symbolic_pymc.utils import graph_equal +from symbolic_pymc.theano.utils import graph_equal def test_terms(): diff --git a/tests/test_linalg.py b/tests/theano/test_linalg.py similarity index 96% rename from tests/test_linalg.py rename to tests/theano/test_linalg.py index f63e4ea..376c9dd 100644 --- a/tests/test_linalg.py +++ b/tests/theano/test_linalg.py @@ -6,9 +6,9 @@ from theano.gof.graph import inputs as tt_inputs -from symbolic_pymc import NormalRV, observed -from symbolic_pymc.meta import mt -from symbolic_pymc.opt import eval_and_reify_meta, FunctionGraph +from symbolic_pymc.theano.random_variables import NormalRV, observed +from symbolic_pymc.theano.meta import mt +from symbolic_pymc.theano.opt import eval_and_reify_meta, FunctionGraph from symbolic_pymc.unify import (etuple, tuple_expression) from symbolic_pymc.relations.linalg import (normal_normal_regression, buildo, normal_qr_transform) diff --git a/tests/test_meta.py b/tests/theano/test_meta.py similarity index 83% rename from tests/test_meta.py rename to tests/theano/test_meta.py index fb2d0c3..4982290 100644 --- a/tests/test_meta.py +++ b/tests/theano/test_meta.py @@ -6,9 +6,12 @@ import pytest from unification import var, isvar, variables -from symbolic_pymc.meta import (MetaSymbol, MetaTensorVariable, MetaTensorType, - mt, metatize) -from symbolic_pymc.utils import graph_equal +from symbolic_pymc.meta import MetaSymbol +from symbolic_pymc.theano.meta import (metatize, + TheanoMetaTensorVariable, + TheanoMetaTensorType, mt) +from symbolic_pymc.theano.utils import graph_equal + def test_metatize(): vec_tt = tt.vector('vec') @@ -50,11 +53,12 @@ class TestOp(tt.gof.Op): assert test_out.obj == TestOp assert test_out.base == TestOp + def test_meta_classes(): vec_tt = tt.vector('vec') vec_m = metatize(vec_tt) assert vec_m.obj == vec_tt - assert type(vec_m) == MetaTensorVariable + assert type(vec_m) == TheanoMetaTensorVariable # This should invalidate the underlying base object. vec_m.index = 0 @@ -63,7 +67,7 @@ def test_meta_classes(): assert vec_m.reify().name == vec_tt.name vec_type_m = vec_m.type - assert type(vec_type_m) == MetaTensorType + assert type(vec_type_m) == TheanoMetaTensorType assert vec_type_m.dtype == vec_tt.dtype assert vec_type_m.broadcastable == vec_tt.type.broadcastable assert vec_type_m.name == vec_tt.type.name @@ -71,7 +75,7 @@ def test_meta_classes(): assert graph_equal(tt.add(1, 2), mt.add(1, 2).reify()) meta_var = mt.add(1, var()).reify() - assert isinstance(meta_var, MetaTensorVariable) + assert isinstance(meta_var, TheanoMetaTensorVariable) assert isinstance(meta_var.owner.op.obj, theano.Op) assert isinstance(meta_var.owner.inputs[0].obj, tt.TensorConstant) diff --git a/tests/test_printing.py b/tests/theano/test_printing.py similarity index 93% rename from tests/test_printing.py rename to tests/theano/test_printing.py index 8ca425a..8173889 100644 --- a/tests/test_printing.py +++ b/tests/theano/test_printing.py @@ -1,7 +1,7 @@ import theano.tensor as tt -from symbolic_pymc import NormalRV -from symbolic_pymc.printing import tt_pprint +from symbolic_pymc.theano.random_variables import NormalRV +from symbolic_pymc.theano.printing import tt_pprint def test_notex_print(): diff --git a/tests/test_pymc3.py b/tests/theano/test_pymc3.py similarity index 85% rename from tests/test_pymc3.py rename to tests/theano/test_pymc3.py index 3a58345..ca4c502 100644 --- a/tests/test_pymc3.py +++ b/tests/theano/test_pymc3.py @@ -13,20 +13,21 @@ # from theano.configparser import change_flags from theano.gof.graph import inputs as tt_inputs -from symbolic_pymc import (MvNormalRV, Observed, observed) -from symbolic_pymc.rv import RandomVariable -from symbolic_pymc.opt import FunctionGraph -from symbolic_pymc.pymc3 import model_graph, graph_model -from symbolic_pymc.utils import canonicalize -from symbolic_pymc.meta import mt +from symbolic_pymc.theano.random_variables import (MvNormalRV, Observed, + observed) +from symbolic_pymc.theano.ops import RandomVariable +from symbolic_pymc.theano.opt import FunctionGraph +from symbolic_pymc.theano.pymc3 import model_graph, graph_model +from symbolic_pymc.theano.utils import canonicalize +from symbolic_pymc.theano.meta import mt def test_pymc_normals(): tt.config.compute_test_value = 'ignore' - mu_X = tt.scalar('mu_X') - sd_X = tt.scalar('sd_X') - mu_Y = tt.scalar('mu_Y') + 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) @@ -34,7 +35,8 @@ def test_pymc_normals(): # 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=0.5) + 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, @@ -48,11 +50,13 @@ def test_pymc_normals(): # This will break comparison if we don't reuse it rng = Z_rv_tt.owner.inputs[1].owner.inputs[-1] - mu_X_ = mt.scalar('mu_X') - sd_X_ = mt.scalar('sd_X') + 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(0., 0.5, None, rng, name='S_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, @@ -60,9 +64,9 @@ def test_pymc_normals(): obs_ = mt(Z_rv.observations) Z_rv_obs_ = mt.observed(obs_, Z_rv_) - Z_rv_meta = canonicalize(Z_rv_obs_.reify(), return_graph=False) + Z_rv_meta = mt(canonicalize(Z_rv_obs_.reify(), return_graph=False)) - assert mt(Z_rv_tt) == mt(Z_rv_meta) + assert mt(Z_rv_tt) == Z_rv_meta # Now, let's try that with multiple outputs. fgraph.disown() diff --git a/tests/test_rv.py b/tests/theano/test_rv.py similarity index 97% rename from tests/test_rv.py rename to tests/theano/test_rv.py index db26c6c..4a055b4 100644 --- a/tests/test_rv.py +++ b/tests/theano/test_rv.py @@ -2,7 +2,7 @@ import theano.tensor as tt -from symbolic_pymc import NormalRV, MvNormalRV +from symbolic_pymc.theano.random_variables import NormalRV, MvNormalRV def rv_numpy_tester(rv, *params, size=None): diff --git a/tests/test_unify.py b/tests/theano/test_unify.py similarity index 96% rename from tests/test_unify.py rename to tests/theano/test_unify.py index 6ff7353..c114f3d 100644 --- a/tests/test_unify.py +++ b/tests/theano/test_unify.py @@ -1,5 +1,4 @@ import pytest -import theano import theano.tensor as tt from operator import add @@ -8,8 +7,8 @@ from kanren.term import term, operator, arguments -from symbolic_pymc.meta import mt -from symbolic_pymc.utils import graph_equal +from symbolic_pymc.theano.meta import mt +from symbolic_pymc.theano.utils import graph_equal from symbolic_pymc.unify import (ExpressionTuple, etuple, tuple_expression) @@ -19,13 +18,9 @@ def test_unification(): y_s = tt.scalar('y_s') c_tt = tt.constant(1, 'c') d_tt = tt.constant(2, 'd') - # x_l = tt.vector('x_l') - # y_l = tt.vector('y_l') - # z_l = tt.vector('z_l') x_l = var('x_l') y_l = var('y_l') - z_l = var('z_l') assert a == reify(x_l, {x_l: a}).reify() test_expr = mt.add(1, mt.mul(2, x_l)) From 6e732cfc956ffc61a609f1e79705b0b52f819ab7 Mon Sep 17 00:00:00 2001 From: "Brandon T. Willard" Date: Fri, 8 Mar 2019 18:39:27 -0600 Subject: [PATCH 2/3] Implement basic TensorFlow meta objects This commit provides basic TensorFlow graph interop. There are still a few design questions and choices to iterate on (e.g. the TF namespace/graph used during intermediate base-object derived meta object construction steps), but the basic class wrappers, helper functions, term representation and unification are working. Closes #3. --- conftest.py | 14 + symbolic_pymc/meta.py | 176 ++++--- symbolic_pymc/tensorflow/__init__.py | 3 + symbolic_pymc/tensorflow/meta.py | 698 ++++++++++++++++++++++++++- symbolic_pymc/tensorflow/unify.py | 9 +- symbolic_pymc/theano/__init__.py | 2 + symbolic_pymc/theano/meta.py | 66 ++- symbolic_pymc/theano/unify.py | 13 +- symbolic_pymc/unify.py | 21 +- tests/tensorflow/__init__.py | 0 tests/tensorflow/test_meta.py | 157 ++++++ tests/tensorflow/test_unify.py | 75 +++ tests/test_etuple.py | 77 +++ tests/theano/test_conjugates.py | 2 + tests/theano/test_kanren.py | 4 + tests/theano/test_linalg.py | 3 + tests/theano/test_meta.py | 2 + tests/theano/test_pymc3.py | 3 + tests/theano/test_relations.py | 86 ++++ tests/theano/test_unify.py | 71 +-- tests/utils.py | 23 + 21 files changed, 1296 insertions(+), 209 deletions(-) create mode 100644 conftest.py create mode 100644 tests/tensorflow/__init__.py create mode 100644 tests/tensorflow/test_meta.py create mode 100644 tests/tensorflow/test_unify.py create mode 100644 tests/test_etuple.py create mode 100644 tests/theano/test_relations.py create mode 100644 tests/utils.py diff --git a/conftest.py b/conftest.py new file mode 100644 index 0000000..b967fea --- /dev/null +++ b/conftest.py @@ -0,0 +1,14 @@ +import pytest + +@pytest.fixture() +def run_with_theano(): + from symbolic_pymc.theano.meta import load_dispatcher + + load_dispatcher() + + +@pytest.fixture() +def run_with_tensorflow(): + from symbolic_pymc.tensorflow.meta import load_dispatcher + + load_dispatcher() diff --git a/symbolic_pymc/meta.py b/symbolic_pymc/meta.py index 451a445..2f303ee 100644 --- a/symbolic_pymc/meta.py +++ b/symbolic_pymc/meta.py @@ -1,13 +1,12 @@ import abc import types -import inspect import reprlib import numpy as np from itertools import chain from functools import partial -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from unification import isvar, Var @@ -39,27 +38,42 @@ def _metatize(obj): return iter([metatize(o) for o in obj]) +def _make_hashable(x): + if isinstance(x, list): + return tuple(x) + elif isinstance(x, np.ndarray): + return x.data.tobytes() + else: + return x + + def _meta_reify_iter(rands): + """Recursively reify an iterable object and return a boolean indicating the presence of un-reifiable objects, if any.""" # We want as many of the rands reified as possible, any_unreified = False reified_rands = [] - for s in rands: + if isinstance(rands, Mapping): + _rands = rands.items() + else: + _rands = rands + + for s in _rands: if isinstance(s, MetaSymbol): rrand = s.reify() - reified_rands += [rrand] + reified_rands.append(rrand) any_unreified |= isinstance(rrand, MetaSymbol) any_unreified |= isvar(rrand) elif MetaSymbol.is_meta(s): - reified_rands += [s] + reified_rands.append(s) any_unreified |= True elif isinstance(s, (list, tuple)): _reified_rands, _any_unreified = _meta_reify_iter(s) - reified_rands += [type(s)(_reified_rands)] + reified_rands.append(type(s)(_reified_rands)) any_unreified |= _any_unreified else: reified_rands += [s] - return reified_rands, any_unreified + return type(rands)(reified_rands), any_unreified class MetaSymbolType(abc.ABCMeta): @@ -75,7 +89,10 @@ def __new__(cls, name, bases, clsdict): def __setattr__(self, attr, obj): """If a slot value is changed, discard any associated non-meta/base objects.""" - if ( + if attr == "obj": + if isinstance(obj, MetaSymbol): + raise ValueError("base object cannot be a meta object!") + elif ( getattr(self, "obj", None) is not None and not isinstance(self.obj, Var) and attr in getattr(self, "__all_slots__", {}) @@ -83,28 +100,24 @@ def __setattr__(self, attr, obj): and getattr(self, attr) != obj ): self.obj = None - elif attr == "obj": - if isinstance(obj, MetaSymbol): - raise ValueError("base object cannot be a meta object!") object.__setattr__(self, attr, obj) clsdict["__setattr__"] = __setattr__ - new_cls = super().__new__(cls, name, bases, clsdict) + @classmethod + def __metatize(cls, obj): + return cls( + *[getattr(obj, s) for s in getattr(cls, "__slots__", [])], + obj=obj + ) - # making sure namespaces are consistent - if isinstance(new_cls.base, type): - if hasattr(new_cls, "_metatize"): - __metatize = new_cls._metatize - else: + clsdict.setdefault("_metatize", __metatize) - def __metatize(obj): - return new_cls( - *[getattr(obj, s) for s in getattr(new_cls, "__slots__", [])], obj=obj - ) + new_cls = super().__new__(cls, name, bases, clsdict) - _metatize.add((new_cls.base,), __metatize) + if isinstance(new_cls.base, type): + _metatize.add((new_cls.base,), new_cls._metatize) # TODO: Could register base classes. # E.g. cls.register(bases) @@ -112,7 +125,10 @@ def __metatize(obj): class MetaSymbol(metaclass=MetaSymbolType): - """Meta objects for unification and such.""" + """Meta objects for unification and such. + + TODO: Should `MetaSymbol.obj` be an abstract property and a `weakref`? + """ @property @abc.abstractmethod @@ -141,7 +157,7 @@ def rands(self): def reify(self): """Create a concrete base object from this meta object (and its rands).""" - if self.obj and not isinstance(self.obj, Var): + if self.obj is not None and not isinstance(self.obj, Var): return self.obj else: reified_rands, any_unreified = _meta_reify_iter(self.rands()) @@ -157,46 +173,30 @@ def reify(self): return res def __eq__(self, other): - """Syntactic equality between meta objects and their bases.""" - # TODO: Allow a sort of cross-inheritance equivalence (e.g. a - # `tt.Variable` or `tt.TensorVariable`)? - # a_sub_b = isinstance(self, type(other)) - # b_sub_a = isinstance(other, type(self)) - # if not (a_sub_b or b_sub_a): - # return False + """Implement an equivalence between meta objects and their bases.""" + if self is other: + return True + if not (type(self) == type(other)): return False - # TODO: ? - # Same for base objects - # a_sub_b = isinstance(self.base, type(other.base)) - # b_sub_a = isinstance(other.base, type(self.base)) - # if not (a_sub_b or b_sub_a): - # return False if not (self.base == other.base): return False - # TODO: ? - # # `self` is the super class, that might be generalizing - # # `other` - a_slots = getattr(self, "__slots__", []) - # b_slots = getattr(other, '__slots__', []) - # if (b_sub_a and not a_sub_b and - # not all(getattr(self, attr) == getattr(other, attr) - # for attr in a_slots)): - # return False - # # `other` is the super class, that might be generalizing - # # `self` - # elif (a_sub_b and not b_sub_a and - # not all(getattr(self, attr) == getattr(other, attr) - # for attr in b_slots)): - # return False - if not all(_check_eq(getattr(self, attr), getattr(other, attr)) for attr in a_slots): + a_slots = getattr(self, "__slots__", None) + if a_slots is not None: + if not all(_check_eq(getattr(self, attr), getattr(other, attr)) for attr in a_slots): + return False + elif getattr(other, "__slots__", None) is not None: + # The other object has slots, but this one doesn't. return False - - # if (self.obj and not isvar(self.obj) and - # other.obj and not isvar(other.objj)): - # assert self.obj == other.obj + else: + # Neither have slots, so best we can do is compare + # base objects (if any). + # If there aren't base objects, we say they're not equal. + # (Maybe we should *require* base objects in this case + # and raise an exception?) + return getattr(self, "obj", None) == getattr(other, "obj", None) is not None return True @@ -204,16 +204,11 @@ def __ne__(self, other): return not self.__eq__(other) def __hash__(self): - def _make_hashable(x): - if isinstance(x, list): - return tuple(x) - elif isinstance(x, np.ndarray): - return x.data.tobytes() - else: - return x - - rands = tuple(_make_hashable(p) for p in self.rands()) - return hash(rands + (self.base,)) + if getattr(self, "__slots__", None) is not None: + rands = tuple(_make_hashable(p) for p in self.rands()) + return hash(rands + (self.base,)) + else: + return hash((self.base, self.obj)) def __str__(self): obj = getattr(self, "obj", None) @@ -246,9 +241,11 @@ def _repr_pretty_(self, p, cycle): p.text(name) p.text("=") p.pretty(item) + obj = getattr(self, "obj", None) - if obj: - if idx: + + if obj is not None: + if idx is not None: p.text(",") p.breakable() p.text("obj=") @@ -261,7 +258,7 @@ def _metatize(obj): class MetaOp(MetaSymbol): - """A meta object that represents a `MetaVaribale`-producing operator. + """A meta object that represents a `MetaVariable`-producing operator. Also, make sure to override `Op.out_meta_type` and make it return the expected meta variable type, if it isn't the default: `MetaTensorVariable`. @@ -274,42 +271,43 @@ class MetaOp(MetaSymbol): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self.op_sig = inspect.signature(self.obj.make_node) @property def obj(self): - return self._obj + return object.__getattribute__(self, "_obj") @obj.setter def obj(self, x): if hasattr(self, "_obj"): raise ValueError("Cannot reset obj in an `Op`") - self._obj = x + object.__setattr__(self, "_obj", x) @abc.abstractmethod - def out_meta_type(self, inputs=None): - """Return the type of meta variable this `Op` is expected to produce given the inputs.""" + def out_meta_types(self, inputs=None): + """Return the types of meta variables this `Op` is expected to produce given the inputs.""" raise NotImplementedError() @abc.abstractmethod def __call__(self, *args, ttype=None, index=None, **kwargs): raise NotImplementedError() - def __eq__(self, other): - # Since these have no rands/slots, we can only really compare against - # the underlying base objects (which should be there!). - if not super().__eq__(other): - return False - - assert self.obj - if self.obj != other.obj: - return False +class MetaVariable(MetaSymbol): + @property + @abc.abstractmethod + def operator(self): + """Return a meta object representing an operator, if any, capable of producing this variable. - return True + It should be callable with all inputs necessary to reproduce this + tensor given by `MetaVariable.inputs`. + """ + raise NotImplementedError() - def __hash__(self): - return hash((self.base, self.obj)) + @property + @abc.abstractmethod + def inputs(self): + """Return the inputs necessary for `MetaVariable.operator` to produced this variable, if any.""" + raise NotImplementedError() def _find_meta_type(obj_type, meta_abs_type): @@ -345,7 +343,3 @@ def _metatize(obj_type): if obj_cls is not None: return obj_cls - - -class MetaVariable(MetaSymbol): - pass diff --git a/symbolic_pymc/tensorflow/__init__.py b/symbolic_pymc/tensorflow/__init__.py index 294e343..ff4d9ed 100644 --- a/symbolic_pymc/tensorflow/__init__.py +++ b/symbolic_pymc/tensorflow/__init__.py @@ -1,3 +1,6 @@ from tensorflow.python.framework import ops +# Needed to register generic functions +from .unify import * + ops.disable_eager_execution() diff --git a/symbolic_pymc/tensorflow/meta.py b/symbolic_pymc/tensorflow/meta.py index e6bfe51..e753c44 100644 --- a/symbolic_pymc/tensorflow/meta.py +++ b/symbolic_pymc/tensorflow/meta.py @@ -1,12 +1,144 @@ +import types +import inspect + import tensorflow as tf +import tensorflow_probability as tfp + +from inspect import Parameter, Signature + +from collections import OrderedDict, UserString + +from functools import partial, wraps + +from unification import Var, var + +from tensorflow.python.framework import (tensor_util, op_def_registry, + op_def_library, tensor_shape) +from tensorflow.core.framework.op_def_pb2 import OpDef + +from tensorflow.core.framework.node_def_pb2 import NodeDef +from tensorflow.core.framework.tensor_shape_pb2 import TensorShapeProto + +from tensorflow_probability import distributions as tfd + + +from ..meta import (MetaSymbol, MetaSymbolType, MetaOp, MetaVariable, + _meta_reify_iter, _metatize, metatize) + + +class MetaOpDefLibrary(op_def_library.OpDefLibrary): + @classmethod + def make_opdef_sig(cls, opdef): + """Create a `Signature` object for an `OpDef`. + + Annotations are include so that one can partially verify arguments. + """ + input_args = OrderedDict([(a.name, a.type or a.type_attr) for a in opdef.input_arg]) + attrs = OrderedDict([(a.name, a.type) for a in opdef.attr]) + + opdef_py_func = getattr(tf.raw_ops, opdef.name, None) + + params = OrderedDict() + if opdef_py_func: + # We assume we're dealing with a function from `tf.raw_ops`. + # Those functions have only the necessary `input_arg`s and + # `attr` inputs as arguments. + opdef_func_sig = Signature.from_callable(opdef_py_func) + + for name, param in opdef_func_sig.parameters.items(): + # We make positional parameters permissible (since the + # functions in `tf.raw_ops` are keyword-only), and we use the + # `tf.raw_ops` arguments to determine the *actual* required + # arguments (because `OpDef`'s `input_arg`s and `attrs` aren't + # exactly clear about that). + if name in input_args: + new_default = Parameter.empty + new_annotation = input_args[name] + else: + new_default = None + new_annotation = attrs.get(name, None) + + new_param = param.replace( + kind=Parameter.POSITIONAL_OR_KEYWORD, + default=new_default, + annotation=new_annotation, + ) + params[name] = new_param + + else: + params = [] + for i_name, i_type in input_args: + p = Parameter(i_name, Parameter.POSITIONAL_OR_KEYWORD, annotation=i_type) + params[i_name] = p + + # These are the ambiguities we're attempting to overcome + # with the `tf.raw_ops` functions above. + for a_name, a_type in attrs: + if a_name == "T": + # This is a type value that will most likely be inferred + # from/by the inputs. + # TODO: We could check for an `allowed_values` attribute. + continue + p = Parameter( + a_name, + Parameter.POSITIONAL_OR_KEYWORD, + # TODO: We could use the `default_value` + # attribute. + default=None, + annotation=a_type, + ) + params[a_name] = p + + # Always assume that a name can be specified. + if "name" not in params: + params["name"] = Parameter("name", Parameter.POSITIONAL_OR_KEYWORD, default=None) -from multipledispatch import dispatch + opdef_sig = Signature( + params.values(), return_annotation=[(o.name, o.type_attr) for o in opdef.output_arg] + ) + return opdef_sig -from ..meta import MetaSymbol, MetaOp, MetaVariable, _metatize + def add_op(self, op_def): + op_info = self._ops.get(op_def.name, None) + if op_info is None: + super().add_op(op_def) + op_info = self._ops[op_def.name] + opdef_sig = self.make_opdef_sig(op_info.op_def) + op_info.opdef_sig = opdef_sig + return op_info -@dispatch(object) -def _metatize(obj): +op_def_lib = MetaOpDefLibrary() + + +class TFlowName(UserString): + """A wrapper for Tensor names. + + TF `Operation` names, and the variables that result from them, cannot be + compared directly due to their uniqueness (within a TF namespace/scope). + + This wrapper class ignores those TF distinctions during string comparison. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._scope_op, _, self._in_idx = self.data.partition(":") + self._scope, _, self._op_name = self._scope_op.partition("/") + self._unique_name = self._op_name.split("_", 1)[0] + + def __eq__(self, other): + if self is other: + return True + + if type(self) != type(other): + return False + elif isinstance(other, str): + return self.data == other + + return self._unique_name == other._unique_name and self._in_idx == other._in_idx + + +def _metatize_tf_object(obj): try: obj = tf.convert_to_tensor(obj) except TypeError: @@ -14,28 +146,564 @@ def _metatize(obj): return _metatize(obj) +def load_dispatcher(): + """Set/override dispatcher to default to TF objects.""" + _metatize.add((object,), _metatize_tf_object) + + +load_dispatcher() + + class TFlowMetaSymbol(MetaSymbol): def reify(self): - # super().reify() - # TODO: Follow `tfp.distribution.Distribution`'s lead? # with tf.name_scope(self.name): # pass - pass + return super().reify() + + +class TFlowMetaOpDef(MetaOp, TFlowMetaSymbol): + """A meta `OpDef`. + + This is like an `Op` node in Theano. + + Some useful info/links: + - https://stackoverflow.com/questions/41147734/looking-for-source-code-of-from-gen-nn-ops-in-tensorflow/41149557#41149557 + + - A better way to view an `OpDef`: + >>> from google.protobuf import json_format + >>> print(json_format.MessageToJson(opdef)) + - If you want to use an `OpDef` to construct a node, see + `op_def_library.OpDefLibrary.apply_op`. + + """ + + base = OpDef + + def __init__(self, obj=None): + op_info = op_def_lib.add_op(obj) + self.apply_func_sig = op_info.opdef_sig + self.apply_func = partial(op_def_lib.apply_op, obj.name) + super().__init__(obj=obj) + + def out_meta_types(self, inputs=None): + out_meta_types = tuple( + # TODO: What other types should we expect? + TFlowMetaTensor if o.type_attr == "T" else None + for i, o in enumerate(self.obj.output_arg) + ) + # TODO: We also have permissible dtype information: + # from objects in the array `self.obj.attr` under the field + # `allowed_values`. + return out_meta_types + + def input_args(self, *args, **kwargs): + kwargs = OrderedDict( + [ + (k, v) + for k, v in kwargs.items() + # Filter out the optional keyword arguments so they we only pass + # expected arguments to the `OpDef`'s apply function. + if k in self.apply_func_sig.parameters + ] + ) + op_args = self.apply_func_sig.bind(*args, **kwargs) + op_args.apply_defaults() + return op_args.arguments + + def __call__(self, *args, **kwargs): + """Create the meta object(s) resulting from an application of this `OpDef`'s implied `Operation`.""" + op_args, op_args_unreified = _meta_reify_iter(args) + op_kwargs, op_kwargs_unreified = _meta_reify_iter(kwargs) + apply_arguments = self.input_args(*op_args, **op_kwargs) + + if not op_args_unreified and not op_kwargs_unreified: + # In this case, we can actually create the TF objects and then turn + # them into meta objects. Doing so will yield information we + # wouldn't be able to produce otherwise (e.g. shape info). + + # TODO: We could make this action/approach configurable (i.e. + # do not perform intermediate/"eager" object construction). + # Especially, when/if we're comfortable with our ability to infer + # the TF-`Operation` inferred values (e.g. shapes, dtypes, etc.) + + # We have to use a primitive string or TF will complain. + name = apply_arguments.get("name", None) + if name is not None: + apply_arguments["name"] = str(name) + + tf_out = self.apply_func(**apply_arguments) + res_var = metatize(tf_out) + return res_var + + # + # If we're here, that means we have to create the meta objects + # manually. + # + + # TODO: `tf.placeholder`s are pretty flexible, we could probably use + # one as a stand-in for any un-reified tensor arguments and at least + # get some partial `dtype`, `shape` and `name` info. + + op_input_args = tuple( + apply_arguments.get(i.name) for i in self.obj.input_arg if i.name in apply_arguments + ) + + # Get the `OpDef`-instantiating parameters and call them a "node_def". + node_def = { + a.name: apply_arguments[a.name] for a in self.obj.attr if a.name in apply_arguments + } + + op_mt = TFlowMetaOp(self, node_def, op_input_args, name=op_kwargs.get("name", None)) + + res_var = op_mt.default_output + + return res_var + + def _repr_pretty_(self, p, cycle): + return p.text(f"{self.__class__.__name__}({self.obj.name})") + + +class TFlowMetaOp(TFlowMetaSymbol): + """A meta `Operation`. + + This is like an `Apply` node in Theano. + + TODO: This whole thing should probably be a "NodeDef" class? + """ -class TFlowMetaOp(MetaOp, TFlowMetaSymbol): base = tf.Operation + __slots__ = ["op_def", "node_def", "inputs", "name"] + + @classmethod + def process_node_def(cls, node_def, op_def, obj=None): + """Minimally convert -Proto objects in a `NodeDef` to their corresponding meta classes. + + The result is a dict that somewhat mirrors a NodeDef, but with meta objects. + + FYI: We're using `node_def` as a hackish way to include user-specified + `OpDef` parameters that aren't standard inputs (e.g. imagine that + `OpDef`s are parameterized and an `OpDef` instance is only possible + with a concrete type parameter). + """ + if isinstance(node_def, NodeDef): + assert node_def.op == op_def.obj.name + assert obj is not None + node_def = dict(node_def.attr) + elif not isinstance(node_def, dict): + raise TypeError("Invalid node_def type") + + op_def_attr_names = tuple( + a.name + for a in op_def.obj.attr + # This is a quick way to filter out the useful + # `attr`s (i.e. the ones that are required `OpDef` + # parameters). + if a.name in op_def.apply_func_sig.parameters + ) + meta_node_def = {} + for k, v in node_def.items(): + if k not in op_def_attr_names: + continue + v = obj.get_attr(k) if obj else v + if k == "shape": + if isinstance(v, TensorShapeProto): + v = tensor_shape.as_shape(v) + v = metatize(v) + elif k == "dtype": + v = tf.as_dtype(v) + meta_node_def[k] = v + + return meta_node_def + + def __init__(self, op_def, node_def, inputs, name=None, outputs=None, obj=None): + if isinstance(op_def, str): + op_def = op_def_registry.get_registered_ops()[op_def] + + self.op_def = metatize(op_def) + + if isinstance(name, (str, TFlowName)) or name is None: + if name is None: + name = op_def.obj.name + # if name and name[-1] == "/": + # name = ops._name_from_scope_name(str(name)) + # else: + # g_tf = ops.get_default_graph() + # name = g_tf.unique_name(str(name)) + self.name = TFlowName(name) + else: + self.name = name + + if self.type is None: + self.type = self.op_def.obj.name + + self.node_def = self.process_node_def(node_def, self.op_def, obj) + self.inputs = tuple(metatize(i) for i in inputs) + + if outputs is not None: + # TODO: Use `weakref`? + self._outputs = tuple(metatize(o) for o in outputs) + else: + self._outputs = None + + super().__init__(obj=obj) + + @property + def type(self): + return self.op_def.obj.name + + @property + def outputs(self): + """Compute meta object outputs for this meta `Operation`. + + If the outputs were specified during construction of this meta + `Operation`, then those outputs are always returned. + + NOTE: These values are dynamically computed, but they could be cached. + One of the reasons that they're dynamically computed: as constituent + meta elements are unified, we may obtain more information about the + """ + if self._outputs is not None: + return self.outputs + + apply_arguments = self.op_def.input_args(*self.inputs, **self.node_def) + out_types_mt = self.op_def.out_meta_types(inputs=apply_arguments) - def __call__(self, *args, ttype=None, index=None, **kwargs): - pass + mt_outs = tuple( + o( + var(), + op=self, + value_index=i, + shape=var(), + name=( + TFlowName(f"{self.name}:{i}") + if isinstance(self.name, (str, TFlowName)) + else var() + ), + ) + for i, o in enumerate(out_types_mt) + ) + return mt_outs -class TFTensorVariable(MetaVariable, TFlowMetaSymbol): + @property + def default_output(self): + """Return the default output for this `Operation`. + + TODO: It might be worth considering a direct approach, and not one that + requires the generation of all meta outputs. + """ + mt_outs = self.outputs + + if len(mt_outs) == 1: + out_var = mt_outs[0] + elif self.value_index is not None and not MetaSymbol.is_meta(self.value_index): + out_var = mt_outs[self.value_index] + else: + # TODO: Should we raise a `ValueError` instead? + out_var = mt_outs + + return out_var + + def reify(self): + if self.obj and not isinstance(self.obj, Var): + return self.obj + + # tt_op = self.op.reify() + # if not self.is_meta(tt_op): + op_inputs, op_inputs_unreified = _meta_reify_iter(self.inputs) + op_attrs, op_attrs_unreified = _meta_reify_iter(self.node_def) + if not op_inputs_unreified and not op_attrs_unreified and not MetaSymbol.is_meta(self.name): + + # We have to use a primitive string or TF will complain. + name = self.name + if self.name is not None: + name = str(name) + + apply_arguments = self.op_def.input_args(*op_inputs, name=name, **op_attrs) + tf_out = self.op_def.apply_func(**apply_arguments) + op_tf = tf_out.op + + assert op_tf is not None + self.obj = op_tf + return self.obj + + return self + + +class TFlowMetaOpFactory(MetaSymbolType): + """A type that enables the creation of meta `OpDef`s by their string names. + + Example + ------- + + >>> TFlowMetaTensor('float64', 'Placeholder') + """ + _op_types = {} + + def __new__(cls, name, bases, clsdict): + op_type = clsdict.get('op_type', None) + new_cls = super().__new__(cls, name, bases, clsdict) + if op_type is not None: + cls._op_types[op_type] = new_cls + return new_cls + + def __call__(cls, dtype=None, op=None, value_index=None, shape=None, name=None, obj=None): + if isinstance(op, str): + # Attempt to construct this meta `Tensor` from the `OpDef` and + # these `Tensor` arguments. + # NOTE: This is really only expected to work for nullary + # `Operation`s (with, perhaps, some optional arguments). + op_def_mt = TFlowMetaOpDef(obj=op_def_registry.get_registered_ops()[op.capitalize()]) + + obj_mt = op_def_mt(dtype=dtype, shape=shape, name=name) + return obj_mt + + return type.__call__( + cls, dtype, op=op, value_index=value_index, shape=shape, name=name, obj=obj + ) + + +class TFlowMetaTensor(MetaVariable, TFlowMetaSymbol, metaclass=TFlowMetaOpFactory): base = tf.Tensor - __slots__ = ["op", "value_index", "dtype"] + __slots__ = ("dtype", "op", "value_index", "shape", "name") - def __init__(self, op, value_index, dtype): - self.op = op - self.value_index = value_index + @classmethod + def _metatize(cls, obj): + """Specialize the meta type based on a `tf.Tensor`'s op.""" + cls = TFlowMetaTensor._op_types.get(obj.op.type, cls) + return cls( + *[getattr(obj, s) for s in getattr(cls, "__slots__", [])], + obj=obj + ) + + def __init__(self, dtype, op=None, value_index=None, shape=None, name=None, obj=None): self.dtype = dtype + self.op = metatize(op) + self.shape = metatize(shape) + self.value_index = value_index + self.name = TFlowName(name) if isinstance(name, str) else name + super().__init__(obj=obj) + + @property + def operator(self): + if self.op is not None: + return self.op.op_def + + @property + def inputs(self): + """Return the tensor's inputs/rands. + + NOTE: These inputs differ from `self.op.inputs` in that contain + the `node_def` parameters, as well. + In other words, these can be used to recreate this object (per + the meta object spec). + """ + if self.op is not None: + input_args = self.op.op_def.input_args( + *self.op.inputs, name=self.op.name, **self.op.node_def + ) + return tuple(input_args.values()) + + def reify(self): + if self.obj is not None and not isinstance(self.obj, Var): + return self.obj + + if not self.op: + op_res = super().reify() + return op_res + + tf_op = self.op.reify() + + if not MetaSymbol.is_meta(tf_op): + + if not MetaSymbol.is_meta(self.value_index): + tf_res = tf_op.outputs[self.value_index] + elif len(tf_op.outputs) == 1: + tf_res = tf_op.outputs[0] + else: + # TODO: Anything else we should/can do here? + return self + + self.obj = tf_res + return tf_res + + return self + + +class TFlowMetaTensorShape(TFlowMetaSymbol): + base = tf.TensorShape + __slots__ = ("dims",) + + def __init__(self, dims, **kwargs): + self.dims = dims + if self.dims is not None: + self.dims = [tensor_shape.as_dimension(d) for d in self.dims] + super().__init__(**kwargs) + + @property + def rank(self): + if self.dims is not None: + return len(self.dims) + + @property + def ndims(self): + return self.rank + + def as_list(self): + return [d.value for d in self.dims] + + +class TFlowConstantType(type): + # def __subclasscheck__(self, c): + + def __instancecheck__(self, o): + if isinstance(o, tf.Tensor): + return o.op.type == "Const" + return False + + +class _TFlowConstant(tf.Tensor, metaclass=TFlowConstantType): + """A helper for `isinstance` functionality.""" + pass + + +class TFlowMetaConstant(TFlowMetaTensor): + base = _TFlowConstant + __slots__ = () + op_type = 'Const' + + def __init__(self, dtype=None, op=None, value_index=None, shape=None, name=None, obj=None): + + assert obj is not None + + if not isinstance(obj, tf.Tensor): + tf_obj = tf.constant(obj, dtype=dtype, shape=shape, name=name) + else: + tf_obj = obj + obj = tensor_util.MakeNdarray(tf_obj.op.get_attr("value")) + + assert tf_obj.op.type == "Const" + self._data = obj + + super().__init__( + tf_obj.dtype.name, + op=tf_obj.op, + value_index=tf_obj.value_index, + name=tf_obj.name, + obj=tf_obj, + ) + + @property + def data(self): + """Return the data for a tensor constant as a Python object. + + TF tensors can also be constants, but there's no separate + class/type for them, so, for access to the underlying constant value, + we provide this property. + """ + if hasattr(self, "_data"): + return self._data + else: + self._data = tensor_util.MakeNdarray(self.op.obj.get_attr("value")) + return self._data + + def __eq__(self, other): + if self is other: + return True + + if not (type(self) == type(other)): + return False + + if not (self.base == other.base): + return False + + our_data = self.op.obj.get_attr("value") + other_data = other.op.obj.get_attr("value") + return our_data == other_data + + def __hash__(self): + data = self.op.obj.get_attr("value") + return hash((self.base, data)) + + +class TFlowMetaAccessor(object): + """An accessor object that simplifies the use of meta objects. + + Instances of this class can be used to implicitly convert TensorFlow + functions and objects into meta objects. + """ + + namespaces = [tf, tf.raw_ops, tfp, tfd] + + def __init__(self, namespace=None): + if namespace is None: + from symbolic_pymc.tensorflow import meta # pylint: disable=import-self + + self.namespaces += [meta] + else: + self.namespaces = [namespace] + + def __call__(self, x): + return metatize(x) + + @classmethod + def find_opdef(cls, obj): + if hasattr(tf.raw_ops, obj.capitalize()): + meta_obj = TFlowMetaOpDef(obj=op_def_registry.get_registered_ops()[obj.capitalize()]) + return meta_obj + return None + + def __getattr__(self, obj): + + ns_obj = next((getattr(ns, obj) for ns in self.namespaces if hasattr(ns, obj)), None) + + if ns_obj is None: + # Try caller's namespace + frame = inspect.currentframe() + f_back = frame.f_back + if f_back: + ns_obj = f_back.f_locals.get(obj, None) + if ns_obj is None: + ns_obj = f_back.f_globals.get(obj) + + if isinstance(ns_obj, (types.FunctionType, partial)): + # We assume that the user requested an `Operation` + # constructor/helper. Return the meta `OpDef`, because + # it implements a constructor/helper-like `__call__`. + meta_obj = self.find_opdef(obj) + + if meta_obj is None: + # It's a function, so let's provide a wrapper that converts + # to-and-from theano and meta objects. + @wraps(ns_obj) + def meta_obj(*args, **kwargs): + args = [o.reify() if hasattr(o, "reify") else o for o in args] + res = ns_obj(*args, **kwargs) + return metatize(res) + + elif isinstance(ns_obj, types.ModuleType): + # It's a sub-module, so let's create another + # `TheanoMetaAccessor` and check within there. + meta_obj = TFlowMetaAccessor(namespace=ns_obj) + else: + + # Hopefully, it's convertible to a meta object... + meta_obj = metatize(ns_obj) + + if meta_obj is None: + # Last resort + meta_obj = self.find_opdef(obj) + + if isinstance(meta_obj, (TFlowMetaSymbol, MetaSymbolType, types.FunctionType)): + setattr(self, obj, meta_obj) + return getattr(self, obj) + elif isinstance(meta_obj, TFlowMetaAccessor): + setattr(self, obj, meta_obj) + return meta_obj + else: + raise AttributeError(f"Meta object for {obj} not found.") + + +mt = TFlowMetaAccessor() diff --git a/symbolic_pymc/tensorflow/unify.py b/symbolic_pymc/tensorflow/unify.py index 08beaef..b3acb58 100644 --- a/symbolic_pymc/tensorflow/unify.py +++ b/symbolic_pymc/tensorflow/unify.py @@ -2,10 +2,10 @@ from kanren.term import term, operator, arguments -from unification.core import reify, _unify, _reify +from unification.core import _reify, _unify, reify from ..meta import metatize -from ..unify import ExpressionTuple, unify_MetaSymbol +from ..unify import ExpressionTuple, unify_MetaSymbol, tuple_expression from .meta import TFlowMetaSymbol tf_class_abstractions = tuple(c.base for c in TFlowMetaSymbol.__subclasses__()) @@ -31,9 +31,12 @@ def _reify_TFlowClasses(o, s): _reify.add((tf_class_abstractions, dict), _reify_TFlowClasses) - operator.add((tf.Tensor,), lambda x: operator(metatize(x))) arguments.add((tf.Tensor,), lambda x: arguments(metatize(x))) term.add((tf.Operation, ExpressionTuple), lambda op, args: term(metatize(op), args)) + +tuple_expression.add(tf_class_abstractions, lambda x: tuple_expression(metatize(x))) + +__all__ = [] diff --git a/symbolic_pymc/theano/__init__.py b/symbolic_pymc/theano/__init__.py index e69de29..bad917f 100644 --- a/symbolic_pymc/theano/__init__.py +++ b/symbolic_pymc/theano/__init__.py @@ -0,0 +1,2 @@ +# Needed to register generic functions +from .unify import * diff --git a/symbolic_pymc/theano/meta.py b/symbolic_pymc/theano/meta.py index c505ba0..146f423 100644 --- a/symbolic_pymc/theano/meta.py +++ b/symbolic_pymc/theano/meta.py @@ -1,6 +1,7 @@ import types import inspect +import weakref import theano import theano.tensor as tt @@ -24,8 +25,7 @@ ) -@dispatch(object) -def _metatize(obj): +def _metatize_theano_object(obj): try: obj = tt.as_tensor_variable(obj) except (ValueError, tt.AsTensorError): @@ -33,6 +33,14 @@ def _metatize(obj): return _metatize(obj) +def load_dispatcher(): + """Set/override dispatcher to default to TF objects.""" + _metatize.add((object,), _metatize_theano_object) + + +load_dispatcher() + + class TheanoMetaSymbol(MetaSymbol): pass @@ -69,22 +77,23 @@ class TheanoMetaOp(MetaOp, TheanoMetaSymbol): `Op.make_node`'s arguments aren't one-to-one with the expected `Apply` node inputs. See `MetaOp.__call__` for more details. - Also, make sure to override `Op.out_meta_type` and make it return the - expected meta variable type, if it isn't the default: `TheanoMetaTensorVariable`. + Also, make sure to override `Op.out_meta_types` and make it return the + expected meta variable types, if it isn't the default: `TheanoMetaTensorVariable`. """ base = tt.Op def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) + self.op_sig = inspect.signature(self.obj.make_node) - def out_meta_type(self, inputs=None): - """Return the type of meta variable this `Op` is expected to produce given the inputs. + def out_meta_types(self, inputs=None): + """Return the types of meta variables this `Op` is expected to produce given the inputs. The default is `TheanoMetaTensorVariable` (corresponding to `TheanoTensorVariable` outputs from the base `Op`). """ - return TheanoMetaTensorVariable + return (TheanoMetaTensorVariable,) def __call__(self, *args, ttype=None, index=None, **kwargs): """Emulate `make_node` for this `Op` and return . @@ -120,12 +129,15 @@ def __call__(self, *args, ttype=None, index=None, **kwargs): res_var = metatize(tt_out) if not isinstance(tt_out, (list, tuple)): - # If the name is indeterminate, we still want all the reified info, - # but we need to make sure that certain parts aren't known. + # If the name is indeterminate, we still want all the reified + # info, but we need to make sure that certain parts aren't + # known. # TODO: In this case, the reified Theano object is a sort of - # "proxy" object; we should use this approach for dtype, as well. - # TODO: We should also put this kind of logic in the appropriate places - # (e.g. `MetaVariable.reify`), when possible. + # "proxy" object; we should use this approach for dtype, as + # well. + # TODO: We should also put this kind of logic in the + # appropriate places (e.g. `MetaVariable.reify`), when + # possible. if TheanoMetaSymbol.is_meta(name): # This should also invalidate `res_var.obj`. res_var.name = name @@ -162,7 +174,7 @@ def __call__(self, *args, ttype=None, index=None, **kwargs): # XXX: We don't have a higher-order meta object model, so being # wrong about the exact type of output variable will cause # problems. - out_meta_type = self.out_meta_type(op_args) + out_meta_type, = self.out_meta_types(op_args) res_var = out_meta_type(ttype, res_apply, index, name) res_var.obj = var() @@ -210,6 +222,11 @@ def __init__(self, op, inputs, outputs=None, obj=None): self.op = metatize(op) self.inputs = tuple(metatize(i) for i in inputs) self.outputs = outputs + if outputs is not None: + # TODO: Convert these to meta objects? + self.outputs = tuple(weakref.ref(o) for o in outputs) + else: + self.outputs = None def reify(self): if self.obj and not isinstance(self.obj, Var): @@ -234,8 +251,6 @@ def nout(self): return len(self.outputs) elif self.obj: return len(self.obj.outputs) - # TODO: Would be cool if we could return - # a logic variable representing this. class TheanoMetaVariable(MetaVariable, TheanoMetaSymbol): @@ -245,10 +260,23 @@ class TheanoMetaVariable(MetaVariable, TheanoMetaSymbol): def __init__(self, type, owner, index, name, obj=None): super().__init__(obj=obj) self.type = metatize(type) - self.owner = metatize(owner) + if owner is not None: + self.owner = metatize(owner) + else: + self.owner = None self.index = index self.name = name + @property + def operator(self): + if self.owner is not None: + return self.owner.op + + @property + def inputs(self): + if self.owner is not None: + return self.owner.inputs + def reify(self): if self.obj and not isinstance(self.obj, Var): return self.obj @@ -300,12 +328,12 @@ class TheanoMetaTensorVariable(TheanoMetaVariable): @property def ndim(self): + # TODO: Would be cool if we could return + # a logic variable representing this. if isinstance(self.type, TheanoMetaTensorType) and isinstance( self.type.broadastable, (list, tuple) ): return len(self.type.broadcastable) - # TODO: Would be cool if we could return - # a logic variable representing this. class TheanoMetaConstant(TheanoMetaVariable): @@ -353,7 +381,7 @@ class TheanoMetaScalarSharedVariable(TheanoMetaSharedVariable): class TheanoMetaAccessor(object): - """Create an object that can be used to implicitly convert Theano functions and object into meta objects. + """Create an object that can be used to implicitly convert Theano functions and objects into meta objects. Use it like a namespace/module/package object, e.g. diff --git a/symbolic_pymc/theano/unify.py b/symbolic_pymc/theano/unify.py index 38a8b90..38a6bc8 100644 --- a/symbolic_pymc/theano/unify.py +++ b/symbolic_pymc/theano/unify.py @@ -1,19 +1,16 @@ import theano.tensor as tt -from multipledispatch import dispatch - from kanren.term import term, operator, arguments from unification.core import _reify, _unify, reify from ..meta import metatize -from ..unify import ExpressionTuple, unify_MetaSymbol -from .meta import mt, TheanoMetaSymbol +from ..unify import ExpressionTuple, unify_MetaSymbol, tuple_expression +from .meta import TheanoMetaSymbol tt_class_abstractions = tuple(c.base for c in TheanoMetaSymbol.__subclasses__()) - _unify.add( (TheanoMetaSymbol, tt_class_abstractions, dict), lambda u, v, s: unify_MetaSymbol(u, metatize(v), s), @@ -41,7 +38,7 @@ def _reify_TheanoClasses(o, s): term.add((tt.Op, ExpressionTuple), lambda op, args: term(metatize(op), args)) +tuple_expression.add(tt_class_abstractions, + lambda x: tuple_expression(metatize(x))) -@dispatch(tt_class_abstractions) -def tuple_expression(x): - return tuple_expression(mt(x)) +__all__ = [] diff --git a/symbolic_pymc/unify.py b/symbolic_pymc/unify.py index 978d11f..f72595b 100644 --- a/symbolic_pymc/unify.py +++ b/symbolic_pymc/unify.py @@ -167,8 +167,17 @@ def unify_MetaSymbol(u, v, s): def _reify_MetaSymbol(o, s): if isinstance(o.obj, Var): + # We allow reification of the base object field for + # a meta object. + # TODO: This is a weird thing that we should probably reconsider. + # It's part of the functionality that allows base objects to fill-in + # as logic variables, though. obj = s.get(o.obj, o.obj) else: + # Otherwise, if there's a base object, it should indicate that there + # are no logic variables or meta terms. + # TODO: Seems like we should be able to skip the reify and comparison + # below. obj = None rands = o.rands() @@ -220,9 +229,9 @@ def operator_MetaVariable(x): `owner.op(owner.inputs)` is consistent, of course. """ - x_owner = getattr(x, "owner", None) - if x_owner and hasattr(x_owner, "op"): - return x_owner.op + x_op = x.operator + if x_op is not None: + return x_op return operator_MetaSymbol(x) @@ -252,9 +261,9 @@ def arguments_MetaVariable(x): `operator_MetaVariable`. """ - x_owner = getattr(x, "owner", None) - if x_owner and hasattr(x_owner, "op"): - x_e = etuple(x_owner.op, *x_owner.inputs, eval_obj=x) + x_op = x.operator + if x_op is not None: + x_e = etuple(x_op, *x.inputs, eval_obj=x) return x_e[1:] return arguments_MetaSymbol(x) diff --git a/tests/tensorflow/__init__.py b/tests/tensorflow/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/tensorflow/test_meta.py b/tests/tensorflow/test_meta.py new file mode 100644 index 0000000..276b987 --- /dev/null +++ b/tests/tensorflow/test_meta.py @@ -0,0 +1,157 @@ +import pytest +import numpy as np + +import tensorflow as tf + +from tensorflow_probability import distributions as tfd + +from symbolic_pymc.tensorflow.meta import (TFlowMetaTensor, + TFlowMetaTensorShape, + TFlowMetaConstant, + TFlowMetaOpDef, + mt) + +from tests.utils import assert_ops_equal + + +@pytest.mark.usefixtures("run_with_tensorflow") +def test_meta_create(): + N = 100 + X = np.vstack([np.random.randn(N), np.ones(N)]).T + X_tf = tf.convert_to_tensor(X) + X_mt = mt(X) + + assert isinstance(X_mt, TFlowMetaTensor) + assert X_mt.op.obj.name.startswith('Const') + # Make sure `reify` returns the cached base object. + assert X_mt.reify() is X_mt.obj + assert isinstance(X_mt.reify(), tf.Tensor) + + assert X_mt == mt(X_tf) + + # from google.protobuf import json_format + # [i for i in X_tf.op.inputs] + # print(json_format.MessageToJson(X_tf.op.op_def)) + + # Create a (constant) tensor meta object manually. + X_raw_mt = TFlowMetaConstant(obj=X) + + assert X_raw_mt._data is X + + # These are *not* equivalent, since they're constants without matching + # constant values (well, our manually-created meta constant has no constant + # value). + assert X_mt == X_raw_mt + # TODO: Should this be true? + # assert X_mt.name == X_raw_mt.name + + add_mt = mt.add(1, 2) + + assert isinstance(add_mt, TFlowMetaTensor) + assert isinstance(add_mt.obj, tf.Tensor) + assert isinstance(add_mt.op.obj, tf.Operation) + assert add_mt.op.obj.type == 'Add' + + assert len(add_mt.op.inputs) == 2 + assert all(isinstance(i, TFlowMetaTensor) + for i in add_mt.op.inputs) + + one_mt, two_mt = mt(1), mt(2) + + assert one_mt != two_mt + + add_mt_2 = mt.add(one_mt, two_mt) + + assert isinstance(add_mt_2, TFlowMetaTensor) + assert isinstance(add_mt_2.obj, tf.Tensor) + assert isinstance(add_mt_2.op.obj, tf.Operation) + assert add_mt_2.op.obj.type == 'Add' + + # These aren't technically equal because of the TF auto-generated names, + # but, since we're using special string wrappers for the names, it should + # work. + assert add_mt == add_mt_2 + + assert add_mt.obj is not None + add_mt.name = None + assert add_mt.obj is None + add_mt_2.name = None + + assert add_mt == add_mt_2 + + a_mt = mt(tf.compat.v1.placeholder('float64', name='a', shape=[1, 2])) + b_mt = mt(tf.compat.v1.placeholder('float64', name='b')) + assert a_mt != b_mt + + assert a_mt.shape.ndims == 2 + assert a_mt.shape == TFlowMetaTensorShape([1, 2]) + + # TODO: Create a placeholder using the string `Operator` name. + z_mt = TFlowMetaTensor('float64', 'Placeholder', name='z__') + + assert z_mt.op.type == 'Placeholder' + assert z_mt.name.startswith('z__') + assert z_mt.obj.name.startswith('z__') + + with pytest.raises(TypeError): + TFlowMetaTensor('float64', 'Add', name='q__') + + # TODO: Test multi-output results + assert True + + +@pytest.mark.usefixtures("run_with_tensorflow") +def test_meta_reify(): + a_mt = mt(tf.compat.v1.placeholder('float64', name='a', shape=[1, 2])) + b_mt = mt(tf.compat.v1.placeholder('float64', name='b', shape=[])) + add_mt = mt.add(a_mt, b_mt) + + assert add_mt.shape.as_list() == [1, 2] + + add_tf = add_mt.reify() + + assert isinstance(add_tf, tf.Tensor) + assert add_tf.op.type == 'Add' + assert add_tf.shape.as_list() == [1, 2] + + # Remove cached base object and force manual reification. + add_mt.obj = None + add_tf = add_mt.reify() + + assert isinstance(add_tf, tf.Tensor) + assert add_tf.op.type == 'Add' + assert add_tf.shape.as_list() == [1, 2] + + +@pytest.mark.usefixtures("run_with_tensorflow") +def test_meta_distributions(): + N = 100 + sigma_tf = tfd.Gamma(np.asarray(1.), np.asarray(1.)).sample() + epsilon_tf = tfd.Normal(np.zeros((N, 1)), sigma_tf).sample() + beta_tf = tfd.Normal(np.zeros((2, 1)), 1).sample() + X = np.vstack([np.random.randn(N), np.ones(N)]).T + X_tf = tf.convert_to_tensor(X) + + Y_tf = tf.linalg.matmul(X_tf, beta_tf) + epsilon_tf + + Y_mt = mt(Y_tf) + + # Confirm that all `Operation`s are the same. + assert_ops_equal(Y_mt, Y_tf) + + # Now, let's see if we can reconstruct it entirely from the + # meta objects. + def _remove_obj(meta_obj): + if (hasattr(meta_obj, 'obj') and + not isinstance(meta_obj, TFlowMetaOpDef)): + meta_obj.obj = None + + if hasattr(meta_obj, 'ancestors'): + for a in meta_obj.ancestors or []: + _remove_obj(a) + + _remove_obj(Y_mt) + + Y_mt_tf = Y_mt.reify() + + assert_ops_equal(Y_mt, Y_mt_tf) diff --git a/tests/tensorflow/test_unify.py b/tests/tensorflow/test_unify.py new file mode 100644 index 0000000..b89771b --- /dev/null +++ b/tests/tensorflow/test_unify.py @@ -0,0 +1,75 @@ +import pytest + +import numpy as np + +import tensorflow as tf + +from unification import unify, reify, var, variables + +from kanren.term import term, operator, arguments + +from symbolic_pymc.tensorflow.meta import (TFlowName, mt) +# from symbolic_pymc.tensorflow.utils import graph_equal +from symbolic_pymc.unify import (ExpressionTuple, etuple, tuple_expression) + +from tests.utils import assert_ops_equal + + +@pytest.mark.usefixtures("run_with_tensorflow") +def test_unification(): + a = tf.compat.v1.placeholder(tf.float64, name='a') + x_l = var('x_l') + a_reif = reify(x_l, {x_l: a}) + assert a_reif.obj is not None + assert a == a_reif.reify() + + test_expr = mt.add(tf.constant(1, dtype=tf.float64), + mt.mul(tf.constant(2, dtype=tf.float64), + x_l)) + test_reify_res = reify(test_expr, {x_l: a}) + test_base_res = test_reify_res.reify() + assert isinstance(test_base_res, tf.Tensor) + + expected_res = (tf.constant(1, dtype=tf.float64) + + tf.constant(2, dtype=tf.float64) * a) + assert_ops_equal(test_base_res, expected_res) + + # from symbolic_pymc.unify import debug_unify; debug_unify() + + meta_expected_res = mt(expected_res) + s_test = unify(test_expr, meta_expected_res, {}) + assert len(s_test) == 5 + + assert reify(test_expr, s_test) == meta_expected_res + + +@pytest.mark.usefixtures("run_with_tensorflow") +def test_etuple_term(): + a = tf.compat.v1.placeholder(tf.float64, name='a') + b = tf.compat.v1.placeholder(tf.float64, name='b') + + a_mt = mt(a) + a_mt.obj = None + a_reified = a_mt.reify() + assert isinstance(a_reified, tf.Tensor) + assert a_reified.shape.dims is None + + test_e = tuple_expression(a_mt) + assert test_e[0] == mt.placeholder + assert test_e[1] == tf.float64 + assert test_e[2][0].base == tf.TensorShape + assert test_e[2][1] is None + + del test_e._eval_obj + a_evaled = test_e.eval_obj + assert all([a == b for a, b in zip(a_evaled.rands(), a_mt.rands())]) + + a_reified = a_evaled.reify() + assert isinstance(a_reified, tf.Tensor) + assert a_reified.shape.dims is None + assert TFlowName(a_reified.name) == TFlowName(a.name) + + e2 = mt.add(a, b) + e2_et = tuple_expression(e2) + assert isinstance(e2_et, ExpressionTuple) + assert e2_et[0] == mt.add diff --git a/tests/test_etuple.py b/tests/test_etuple.py new file mode 100644 index 0000000..752ed51 --- /dev/null +++ b/tests/test_etuple.py @@ -0,0 +1,77 @@ +import pytest + +from operator import add + +from kanren.term import term, operator, arguments + +from symbolic_pymc.unify import (ExpressionTuple, etuple) + + +def test_etuple(): + """Test basic `etuple` functionality. + """ + def test_op(*args): + return tuple(object() for i in range(sum(args))) + + e1 = etuple(test_op, 1, 2) + + assert not hasattr(e1, '_eval_obj') + + with pytest.raises(ValueError): + e1.eval_obj = 1 + + e1_obj = e1.eval_obj + assert len(e1_obj) == 3 + assert all(type(o) == object for o in e1_obj) + + # Make sure we don't re-create the cached `eval_obj` + e1_obj_2 = e1.eval_obj + assert e1_obj == e1_obj_2 + + # Confirm that evaluation is recursive + e2 = etuple(add, (object(),), e1) + + # Make sure we didn't convert this single tuple value to + # an `etuple` + assert type(e2[1]) == tuple + + # Slices should be `etuple`s, though. + assert isinstance(e2[:1], ExpressionTuple) + assert e2[1] == e2[1:2][0] + + e2_obj = e2.eval_obj + + assert type(e2_obj) == tuple + assert len(e2_obj) == 4 + assert all(type(o) == object for o in e2_obj) + # Make sure that it used `e1`'s original `eval_obj` + assert e2_obj[1:] == e1_obj + + # Confirm that any combination of `tuple`s/`etuple`s in + # concatenation result in an `etuple` + e_radd = (1,) + etuple(2, 3) + assert isinstance(e_radd, ExpressionTuple) + assert e_radd == (1, 2, 3) + + e_ladd = etuple(1, 2) + (3,) + assert isinstance(e_ladd, ExpressionTuple) + assert e_ladd == (1, 2, 3) + + +def test_etuple_term(): + """Test `tuple_expression` and `etuple` interaction with `term` + """ + # Make sure that we don't lose underlying `eval_obj`s + # when taking apart and re-creating expression tuples + # using `kanren`'s `operator`, `arguments` and `term` + # functions. + e1 = etuple(add, (object(),), (object(),)) + e1_obj = e1.eval_obj + + e1_dup = (operator(e1),) + arguments(e1) + + assert isinstance(e1_dup, ExpressionTuple) + assert e1_dup.eval_obj == e1_obj + + e1_dup_2 = term(operator(e1), arguments(e1)) + assert e1_dup_2 == e1_obj diff --git a/tests/theano/test_conjugates.py b/tests/theano/test_conjugates.py index d801baa..d1af4a3 100644 --- a/tests/theano/test_conjugates.py +++ b/tests/theano/test_conjugates.py @@ -1,3 +1,4 @@ +import pytest import theano.tensor as tt import numpy as np @@ -10,6 +11,7 @@ from symbolic_pymc.relations.conjugates import conjugate_posteriors +@pytest.mark.usefixtures("run_with_theano") def test_mvnormal_mvnormal(): """Test that we can produce the closed-form distribution for the conjugate multivariate normal-regression with normal-prior model. diff --git a/tests/theano/test_kanren.py b/tests/theano/test_kanren.py index e19656a..575b490 100644 --- a/tests/theano/test_kanren.py +++ b/tests/theano/test_kanren.py @@ -1,3 +1,4 @@ +import pytest import theano.tensor as tt from kanren import run, eq, variables @@ -11,6 +12,7 @@ from symbolic_pymc.theano.utils import graph_equal +@pytest.mark.usefixtures("run_with_theano") def test_terms(): x, a, b = tt.dvectors('xab') test_expr = x + a * b @@ -34,6 +36,7 @@ def test_terms(): arguments(test_expr)).reify() +@pytest.mark.usefixtures("run_with_theano") def test_kanren(): # x, a, b = tt.dvectors('xab') # @@ -78,6 +81,7 @@ def test_kanren(): assert res.reify() == Y_rv +@pytest.mark.usefixtures("run_with_theano") def test_assoccomm(): from kanren.assoccomm import buildo diff --git a/tests/theano/test_linalg.py b/tests/theano/test_linalg.py index 376c9dd..47483f4 100644 --- a/tests/theano/test_linalg.py +++ b/tests/theano/test_linalg.py @@ -1,4 +1,5 @@ import pytest + import theano import theano.tensor as tt @@ -16,6 +17,7 @@ from kanren import run, eq, var +@pytest.mark.usefixtures("run_with_theano") @pytest.mark.xfail(strict=True) def test_normal_normal_regression(): tt.config.compute_test_value = 'ignore' @@ -93,6 +95,7 @@ def test_normal_normal_regression(): assert Y_out_mt == Y_new_mt +@pytest.mark.usefixtures("run_with_theano") @pytest.mark.xfail(strict=True) def test_normal_qr_transform(): np.random.seed(9283) diff --git a/tests/theano/test_meta.py b/tests/theano/test_meta.py index 4982290..b0dab16 100644 --- a/tests/theano/test_meta.py +++ b/tests/theano/test_meta.py @@ -13,6 +13,7 @@ from symbolic_pymc.theano.utils import graph_equal +@pytest.mark.usefixtures("run_with_theano") def test_metatize(): vec_tt = tt.vector('vec') vec_m = metatize(vec_tt) @@ -54,6 +55,7 @@ class TestOp(tt.gof.Op): assert test_out.base == TestOp +@pytest.mark.usefixtures("run_with_theano") def test_meta_classes(): vec_tt = tt.vector('vec') vec_m = metatize(vec_tt) diff --git a/tests/theano/test_pymc3.py b/tests/theano/test_pymc3.py index ca4c502..7f91f6d 100644 --- a/tests/theano/test_pymc3.py +++ b/tests/theano/test_pymc3.py @@ -22,6 +22,7 @@ from symbolic_pymc.theano.meta import mt +@pytest.mark.usefixtures("run_with_theano") def test_pymc_normals(): tt.config.compute_test_value = 'ignore' @@ -104,6 +105,7 @@ def test_pymc_normals(): assert all(v == 1 for v in Z_vars_count.values()) +@pytest.mark.usefixtures("run_with_theano") def test_normals_to_model(): tt.config.compute_test_value = 'ignore' @@ -147,6 +149,7 @@ def test_normals_to_model(): assert isinstance(beta_pm, pm.MvNormal) +@pytest.mark.usefixtures("run_with_theano") def test_pymc_broadcastable(): tt.config.compute_test_value = 'ignore' diff --git a/tests/theano/test_relations.py b/tests/theano/test_relations.py new file mode 100644 index 0000000..80f8fd5 --- /dev/null +++ b/tests/theano/test_relations.py @@ -0,0 +1,86 @@ +import pytest + +import numpy as np +import theano +import theano.tensor as tt + +from theano.gof.opt import EquilibriumOptimizer + +from symbolic_pymc.theano.random_variables import (observed, NormalRV, + HalfCauchyRV) +from symbolic_pymc.theano.opt import KanrenRelationSub, FunctionGraph +from symbolic_pymc.theano.utils import (optimize_graph, canonicalize, + get_rv_observation) +from symbolic_pymc.relations.distributions import scale_loc_transform + + +@pytest.mark.usefixtures("run_with_theano") +def test_pymc_normals(): + tt.config.compute_test_value = 'ignore' + + rand_state = theano.shared(np.random.RandomState()) + mu_a = NormalRV(0., 100**2, name='mu_a', rng=rand_state) + sigma_a = HalfCauchyRV(5, name='sigma_a', rng=rand_state) + mu_b = NormalRV(0., 100**2, name='mu_b', rng=rand_state) + sigma_b = HalfCauchyRV(5, name='sigma_b', rng=rand_state) + county_idx = np.r_[1, 1, 2, 3] + # We want the following for a, b: + # N(m, S) -> m + N(0, 1) * S + a = NormalRV(mu_a, sigma_a, size=(len(county_idx),), name='a', rng=rand_state) + b = NormalRV(mu_b, sigma_b, size=(len(county_idx),), name='b', rng=rand_state) + radon_est = a[county_idx] + b[county_idx] * 7 + eps = HalfCauchyRV(5, name='eps', rng=rand_state) + radon_like = NormalRV(radon_est, eps, name='radon_like', rng=rand_state) + radon_like_rv = observed(tt.as_tensor_variable(np.r_[1., 2., 3., 4.]), radon_like) + + inputs = [mu_a, mu_b, eps, rand_state] + fgraph = FunctionGraph( + inputs, + [radon_like_rv], + clone=True) + fgraph = canonicalize(fgraph, return_graph=True, in_place=False) + + posterior_opt = EquilibriumOptimizer( + [KanrenRelationSub(scale_loc_transform, + node_filter=get_rv_observation)], + max_use_ratio=10) + + fgraph_opt = optimize_graph(fgraph, posterior_opt, return_graph=True) + fgraph_opt = canonicalize(fgraph_opt, return_graph=True, in_place=False) + + radon_like_rv_opt = fgraph_opt.outputs[0] + radon_like_opt = radon_like_rv_opt.owner.inputs[1] + radon_est_opt = radon_like_opt.owner.inputs[0] + + # These should now be `tt.add(mu_*, ...)` outputs. + a_opt = radon_est_opt.owner.inputs[0].owner.inputs[0] + b_opt = radon_est_opt.owner.inputs[1].owner.inputs[1].owner.inputs[0] + + # Make sure NormalRV gets replaced with an addition + assert a_opt.owner.op == tt.add + assert b_opt.owner.op == tt.add + + # Make sure the first term in the addition is the old NormalRV mean + mu_a_opt = a_opt.owner.inputs[0].owner.inputs[0] + assert 'mu_a' == mu_a_opt.name == mu_a.name + mu_b_opt = b_opt.owner.inputs[0].owner.inputs[0] + assert 'mu_b' == mu_b_opt.name == mu_b.name + + # Make sure the second term in the addition is the standard NormalRV times + # the old std. dev. + assert a_opt.owner.inputs[1].owner.op == tt.mul + assert b_opt.owner.inputs[1].owner.op == tt.mul + + sigma_a_opt = a_opt.owner.inputs[1].owner.inputs[0].owner.inputs[0] + assert sigma_a_opt.owner.op == sigma_a.owner.op + sigma_b_opt = b_opt.owner.inputs[1].owner.inputs[0].owner.inputs[0] + assert sigma_b_opt.owner.op == sigma_b.owner.op + + a_std_norm_opt = a_opt.owner.inputs[1].owner.inputs[1] + assert a_std_norm_opt.owner.op == NormalRV + assert a_std_norm_opt.owner.inputs[0].data == 0.0 + assert a_std_norm_opt.owner.inputs[1].data == 1.0 + b_std_norm_opt = b_opt.owner.inputs[1].owner.inputs[1] + assert b_std_norm_opt.owner.op == NormalRV + assert b_std_norm_opt.owner.inputs[0].data == 0.0 + assert b_std_norm_opt.owner.inputs[1].data == 1.0 diff --git a/tests/theano/test_unify.py b/tests/theano/test_unify.py index c114f3d..2f26083 100644 --- a/tests/theano/test_unify.py +++ b/tests/theano/test_unify.py @@ -1,4 +1,5 @@ import pytest + import theano.tensor as tt from operator import add @@ -12,6 +13,7 @@ from symbolic_pymc.unify import (ExpressionTuple, etuple, tuple_expression) +@pytest.mark.usefixtures("run_with_theano") def test_unification(): x, y, a, b = tt.dvectors('xyab') x_s = tt.scalar('x_s') @@ -75,75 +77,10 @@ def test_unification(): unify(mt_expr_add, tt_expr_add_3)).reify()) -def test_etuple(): - """Test basic `etuple` functionality. - """ - def test_op(*args): - return tuple(object() for i in range(sum(args))) - - e1 = etuple(test_op, 1, 2) - - assert not hasattr(e1, '_eval_obj') - - with pytest.raises(ValueError): - e1.eval_obj = 1 - - e1_obj = e1.eval_obj - assert len(e1_obj) == 3 - assert all(type(o) == object for o in e1_obj) - - # Make sure we don't re-create the cached `eval_obj` - e1_obj_2 = e1.eval_obj - assert e1_obj == e1_obj_2 - - # Confirm that evaluation is recursive - e2 = etuple(add, (object(),), e1) - - # Make sure we didn't convert this single tuple value to - # an `etuple` - assert type(e2[1]) == tuple - - # Slices should be `etuple`s, though. - assert isinstance(e2[:1], ExpressionTuple) - assert e2[1] == e2[1:2][0] - - e2_obj = e2.eval_obj - - assert type(e2_obj) == tuple - assert len(e2_obj) == 4 - assert all(type(o) == object for o in e2_obj) - # Make sure that it used `e1`'s original `eval_obj` - assert e2_obj[1:] == e1_obj - - # Confirm that any combination of `tuple`s/`etuple`s in - # concatenation result in an `etuple` - e_radd = (1,) + etuple(2, 3) - assert isinstance(e_radd, ExpressionTuple) - assert e_radd == (1, 2, 3) - - e_ladd = etuple(1, 2) + (3,) - assert isinstance(e_ladd, ExpressionTuple) - assert e_ladd == (1, 2, 3) - - +@pytest.mark.usefixtures("run_with_theano") def test_etuple_term(): """Test `tuple_expression` and `etuple` interaction with `term` """ - # Make sure that we don't lose underlying `eval_obj`s - # when taking apart and re-creating expression tuples - # using `kanren`'s `operator`, `arguments` and `term` - # functions. - e1 = etuple(add, (object(),), (object(),)) - e1_obj = e1.eval_obj - - e1_dup = (operator(e1),) + arguments(e1) - - assert isinstance(e1_dup, ExpressionTuple) - assert e1_dup.eval_obj == e1_obj - - e1_dup_2 = term(operator(e1), arguments(e1)) - assert e1_dup_2 == e1_obj - # Take apart an already constructed/evaluated meta # object. e2 = mt.add(mt.vector(), mt.vector()) @@ -182,4 +119,4 @@ def test_etuple_term(): e2_et_2 = tuple_expression(tt_expr) assert e2_et_2 == e3 == e2_et assert isinstance(e2_et_2, ExpressionTuple) - assert e2_et_2.eval_obj.reify() == tt_expr + assert e2_et_2.eval_obj == tt_expr diff --git a/tests/utils.py b/tests/utils.py new file mode 100644 index 0000000..8b73049 --- /dev/null +++ b/tests/utils.py @@ -0,0 +1,23 @@ +from collections.abc import Mapping + + +def assert_ops_equal(a, b, compare_fn=lambda a, b: a.op.type == b.op.type): + if hasattr(a, 'op') or hasattr(b, 'op'): + assert hasattr(a, 'op') and hasattr(b, 'op') + + assert compare_fn(a, b) + + assert len(a.op.inputs) == len(b.op.inputs) + + if isinstance(a.op, Mapping): + a_inputs = list(a.op.inputs.values()) + else: + a_inputs = list(a.op.inputs) + + if isinstance(b.op, Mapping): + b_inputs = list(b.op.inputs.values()) + else: + b_inputs = list(b.op.inputs) + + for i_a, i_b in zip(a_inputs, b_inputs): + assert_ops_equal(i_a, i_b) From 6daa63de2b3c6e1e004bff2351d4e2b2e2190268 Mon Sep 17 00:00:00 2001 From: "Brandon T. Willard" Date: Thu, 30 May 2019 16:54:56 -0500 Subject: [PATCH 3/3] Move backend-specific relations to their own directories --- symbolic_pymc/meta.py | 5 +-- symbolic_pymc/relations/__init__.py | 29 ----------------- .../relations/tensorflow/__init__.py | 0 symbolic_pymc/relations/theano/__init__.py | 0 .../relations/{ => theano}/conjugates.py | 6 ++-- .../relations/{ => theano}/distributions.py | 31 +++++++++++++++++-- .../relations/{ => theano}/linalg.py | 4 +-- symbolic_pymc/tensorflow/meta.py | 25 +++++++++------ symbolic_pymc/theano/meta.py | 1 - symbolic_pymc/theano/pymc3.py | 6 ++-- symbolic_pymc/theano/unify.py | 3 +- tests/theano/test_conjugates.py | 2 +- tests/theano/test_linalg.py | 4 +-- tests/theano/test_relations.py | 2 +- 14 files changed, 56 insertions(+), 62 deletions(-) create mode 100644 symbolic_pymc/relations/tensorflow/__init__.py create mode 100644 symbolic_pymc/relations/theano/__init__.py rename symbolic_pymc/relations/{ => theano}/conjugates.py (98%) rename symbolic_pymc/relations/{ => theano}/distributions.py (88%) rename symbolic_pymc/relations/{ => theano}/linalg.py (97%) diff --git a/symbolic_pymc/meta.py b/symbolic_pymc/meta.py index 2f303ee..ecb5cbb 100644 --- a/symbolic_pymc/meta.py +++ b/symbolic_pymc/meta.py @@ -107,10 +107,7 @@ def __setattr__(self, attr, obj): @classmethod def __metatize(cls, obj): - return cls( - *[getattr(obj, s) for s in getattr(cls, "__slots__", [])], - obj=obj - ) + return cls(*[getattr(obj, s) for s in getattr(cls, "__slots__", [])], obj=obj) clsdict.setdefault("_metatize", __metatize) diff --git a/symbolic_pymc/relations/__init__.py b/symbolic_pymc/relations/__init__.py index e616d70..1e83f92 100644 --- a/symbolic_pymc/relations/__init__.py +++ b/symbolic_pymc/relations/__init__.py @@ -1,13 +1,6 @@ -import numpy as np -import theano.tensor as tt - from kanren.facts import Relation from kanren.goals import goalify -from unification.utils import transitive_get as walk - -from ..theano.meta import TheanoMetaConstant - # Hierarchical models that we recognize. hierarchical_model = Relation("hierarchical") @@ -17,25 +10,3 @@ concat = goalify(lambda *args: "".join(args)) - - -def constant_neq(lvar, val): - """ - Assert that a constant graph variable is not equal to a specific value. - - Scalar values are broadcast across arrays. - - """ - - def _goal(s): - lvar_val = walk(lvar, s) - if isinstance(lvar_val, (tt.Constant, TheanoMetaConstant)): - 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 - - return _goal diff --git a/symbolic_pymc/relations/tensorflow/__init__.py b/symbolic_pymc/relations/tensorflow/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/symbolic_pymc/relations/theano/__init__.py b/symbolic_pymc/relations/theano/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/symbolic_pymc/relations/conjugates.py b/symbolic_pymc/relations/theano/conjugates.py similarity index 98% rename from symbolic_pymc/relations/conjugates.py rename to symbolic_pymc/relations/theano/conjugates.py index 8ba3e15..9503929 100644 --- a/symbolic_pymc/relations/conjugates.py +++ b/symbolic_pymc/relations/theano/conjugates.py @@ -4,9 +4,9 @@ from kanren import conde, eq from kanren.facts import fact -from . import conjugate -from ..unify import etuple -from ..theano.meta import mt +from .. import conjugate +from ...unify import etuple +from ...theano.meta import mt mt.namespaces += [theano.tensor.nlinalg] diff --git a/symbolic_pymc/relations/distributions.py b/symbolic_pymc/relations/theano/distributions.py similarity index 88% rename from symbolic_pymc/relations/distributions.py rename to symbolic_pymc/relations/theano/distributions.py index 87c1899..890f79e 100644 --- a/symbolic_pymc/relations/distributions.py +++ b/symbolic_pymc/relations/theano/distributions.py @@ -1,11 +1,16 @@ """Relations pertaining to probability distributions.""" +import numpy as np +import theano.tensor as tt + from unification import var from kanren import conde, eq from kanren.facts import fact -from . import constant_neq, concat -from ..unify import etuple -from ..theano.meta import mt +from unification.utils import transitive_get as walk + +from .. import concat +from ...unify import etuple +from ...theano.meta import mt, TheanoMetaConstant from kanren.facts import Relation @@ -50,6 +55,26 @@ # None) +def constant_neq(lvar, val): + """Assert that a constant graph variable is not equal to a specific value. + + Scalar values are broadcast across arrays. + """ + + def _goal(s): + lvar_val = walk(lvar, s) + if isinstance(lvar_val, (tt.Constant, TheanoMetaConstant)): + 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 + + return _goal + + def scale_loc_transform(in_expr, out_expr): """Create relations for lifting and sinking scale and location parameters of distributions. diff --git a/symbolic_pymc/relations/linalg.py b/symbolic_pymc/relations/theano/linalg.py similarity index 97% rename from symbolic_pymc/relations/linalg.py rename to symbolic_pymc/relations/theano/linalg.py index f729e3a..e2e9ba5 100644 --- a/symbolic_pymc/relations/linalg.py +++ b/symbolic_pymc/relations/theano/linalg.py @@ -11,8 +11,8 @@ from kanren.goals import not_equalo, conso from kanren.term import term, operator, arguments -from ..theano.meta import mt -from ..unify import etuple, tuple_expression, ExpressionTuple +from ...theano.meta import mt +from ...unify import etuple, tuple_expression, ExpressionTuple mt.nlinalg.qr_full = mt(QRFull("reduced")) diff --git a/symbolic_pymc/tensorflow/meta.py b/symbolic_pymc/tensorflow/meta.py index e753c44..e56b2d4 100644 --- a/symbolic_pymc/tensorflow/meta.py +++ b/symbolic_pymc/tensorflow/meta.py @@ -12,8 +12,7 @@ from unification import Var, var -from tensorflow.python.framework import (tensor_util, op_def_registry, - op_def_library, tensor_shape) +from tensorflow.python.framework import tensor_util, op_def_registry, op_def_library, tensor_shape from tensorflow.core.framework.op_def_pb2 import OpDef from tensorflow.core.framework.node_def_pb2 import NodeDef @@ -22,8 +21,15 @@ from tensorflow_probability import distributions as tfd -from ..meta import (MetaSymbol, MetaSymbolType, MetaOp, MetaVariable, - _meta_reify_iter, _metatize, metatize) +from ..meta import ( + MetaSymbol, + MetaSymbolType, + MetaOp, + MetaVariable, + _meta_reify_iter, + _metatize, + metatize, +) class MetaOpDefLibrary(op_def_library.OpDefLibrary): @@ -440,10 +446,11 @@ class TFlowMetaOpFactory(MetaSymbolType): >>> TFlowMetaTensor('float64', 'Placeholder') """ + _op_types = {} def __new__(cls, name, bases, clsdict): - op_type = clsdict.get('op_type', None) + op_type = clsdict.get("op_type", None) new_cls = super().__new__(cls, name, bases, clsdict) if op_type is not None: cls._op_types[op_type] = new_cls @@ -473,10 +480,7 @@ class TFlowMetaTensor(MetaVariable, TFlowMetaSymbol, metaclass=TFlowMetaOpFactor def _metatize(cls, obj): """Specialize the meta type based on a `tf.Tensor`'s op.""" cls = TFlowMetaTensor._op_types.get(obj.op.type, cls) - return cls( - *[getattr(obj, s) for s in getattr(cls, "__slots__", [])], - obj=obj - ) + return cls(*[getattr(obj, s) for s in getattr(cls, "__slots__", [])], obj=obj) def __init__(self, dtype, op=None, value_index=None, shape=None, name=None, obj=None): self.dtype = dtype @@ -566,13 +570,14 @@ def __instancecheck__(self, o): class _TFlowConstant(tf.Tensor, metaclass=TFlowConstantType): """A helper for `isinstance` functionality.""" + pass class TFlowMetaConstant(TFlowMetaTensor): base = _TFlowConstant __slots__ = () - op_type = 'Const' + op_type = "Const" def __init__(self, dtype=None, op=None, value_index=None, shape=None, name=None, obj=None): diff --git a/symbolic_pymc/theano/meta.py b/symbolic_pymc/theano/meta.py index 146f423..2aa5c09 100644 --- a/symbolic_pymc/theano/meta.py +++ b/symbolic_pymc/theano/meta.py @@ -8,7 +8,6 @@ from functools import partial, wraps from unification import var, isvar, Var -from multipledispatch import dispatch from kanren.facts import fact from kanren.assoccomm import commutative, associative diff --git a/symbolic_pymc/theano/pymc3.py b/symbolic_pymc/theano/pymc3.py index 54d42f6..e7ca185 100644 --- a/symbolic_pymc/theano/pymc3.py +++ b/symbolic_pymc/theano/pymc3.py @@ -96,8 +96,7 @@ def _convert_rv_to_dist(op, rv): @dispatch(pm.HalfNormal, object) def convert_dist_to_rv(dist, rng): size = dist.shape.astype(int)[HalfNormalRV.ndim_supp :] - res = HalfNormalRV(np.array(0.0, dtype=dist.dtype), - dist.sd, size=size, rng=rng) + res = HalfNormalRV(np.array(0.0, dtype=dist.dtype), dist.sd, size=size, rng=rng) return res @@ -176,8 +175,7 @@ def _convert_rv_to_dist(op, rv): @dispatch(pm.HalfCauchy, object) def convert_dist_to_rv(dist, rng): size = dist.shape.astype(int)[HalfCauchyRV.ndim_supp :] - res = HalfCauchyRV(np.array(0.0, dtype=dist.dtype), - dist.beta, size=size, rng=rng) + res = HalfCauchyRV(np.array(0.0, dtype=dist.dtype), dist.beta, size=size, rng=rng) return res diff --git a/symbolic_pymc/theano/unify.py b/symbolic_pymc/theano/unify.py index 38a6bc8..3be8d31 100644 --- a/symbolic_pymc/theano/unify.py +++ b/symbolic_pymc/theano/unify.py @@ -38,7 +38,6 @@ def _reify_TheanoClasses(o, s): term.add((tt.Op, ExpressionTuple), lambda op, args: term(metatize(op), args)) -tuple_expression.add(tt_class_abstractions, - lambda x: tuple_expression(metatize(x))) +tuple_expression.add(tt_class_abstractions, lambda x: tuple_expression(metatize(x))) __all__ = [] diff --git a/tests/theano/test_conjugates.py b/tests/theano/test_conjugates.py index d1af4a3..529fbac 100644 --- a/tests/theano/test_conjugates.py +++ b/tests/theano/test_conjugates.py @@ -8,7 +8,7 @@ from symbolic_pymc.theano.random_variables import MvNormalRV, observed from symbolic_pymc.theano.opt import KanrenRelationSub, FunctionGraph from symbolic_pymc.theano.utils import optimize_graph -from symbolic_pymc.relations.conjugates import conjugate_posteriors +from symbolic_pymc.relations.theano.conjugates import conjugate_posteriors @pytest.mark.usefixtures("run_with_theano") diff --git a/tests/theano/test_linalg.py b/tests/theano/test_linalg.py index 47483f4..e1acbc5 100644 --- a/tests/theano/test_linalg.py +++ b/tests/theano/test_linalg.py @@ -11,8 +11,8 @@ from symbolic_pymc.theano.meta import mt from symbolic_pymc.theano.opt import eval_and_reify_meta, FunctionGraph from symbolic_pymc.unify import (etuple, tuple_expression) -from symbolic_pymc.relations.linalg import (normal_normal_regression, buildo, - normal_qr_transform) +from symbolic_pymc.relations.theano.linalg import (normal_normal_regression, buildo, + normal_qr_transform) from kanren import run, eq, var diff --git a/tests/theano/test_relations.py b/tests/theano/test_relations.py index 80f8fd5..279698b 100644 --- a/tests/theano/test_relations.py +++ b/tests/theano/test_relations.py @@ -11,7 +11,7 @@ from symbolic_pymc.theano.opt import KanrenRelationSub, FunctionGraph from symbolic_pymc.theano.utils import (optimize_graph, canonicalize, get_rv_observation) -from symbolic_pymc.relations.distributions import scale_loc_transform +from symbolic_pymc.relations.theano.distributions import scale_loc_transform @pytest.mark.usefixtures("run_with_theano")