Commit 61d9a1e2 authored by Michael Green's avatar Michael Green 🚀
Browse files

Merge branch 'redesign-rv-pair' into 'master' 

remove special handeling of rvpair

Closes #33

See merge request !49
parents c65fa60b 7a565f71

setup.py

+1 −4
Original line number Diff line number Diff line
@@ -116,11 +116,8 @@ EXTRAS_REQUIRE = { 
        # graph neural networks tutorial

        "torch",

        "torchvision",

        # "torch-scatter",

        "torch-sparse",

        # "torch-cluster"

        # "torch-spline-conv",

        "torch-geometric"

        "torch-geometric",

    ],

    "examples": ["notebook"],

    "lint": ["black==20.8b1", "isort==4.3.21", "pylint==2.4.4"],

src/borch/distributions/rv_distributions.py

+1 −1
Original line number Diff line number Diff line
@@ -383,7 +383,7 @@ class TransformedDistribution(RandomVariable): 


    def _distribution(self):

        return self.distribution_cls(

            self.base_distribution.distribution,

            self.get("base_distribution").distribution,

            self.transforms,

            validate_args=self.validate_args,

        )

src/borch/module.py

+8 −9
Original line number Diff line number Diff line
@@ -456,10 +456,13 @@ class Module(_Module): 
                if rv is not None:

                    rv()  # redraw this sample



    def get(self, name):

        """Standard `getattr` with no custom overloading"""

        return super().__getattr__(name)



    def __getattr__(self, name):

        try:

            out = super().__getattr__(name)

            return out

            return super().__getattr__(name)

        except AttributeError:

            posterior = self.__dict__.get("_modules", {}).get("posterior", None)

            prior = self.__dict__.get("_modules", {}).get("prior", None)
@@ -468,7 +471,7 @@ class Module(_Module): 
                if param is None and isinstance(

                    getattr(prior, name, None), borch.RandomVariable

                ):

                    self.posterior.set_random_variable(name, getattr(prior, name), None)

                    self.posterior.set_random_variable(name, getattr(prior, name))

                    param = getattr(posterior, name)

                if isinstance(param, borch.RandomVariable):

                    observed = self.observed.get(name)
@@ -484,19 +487,15 @@ class Module(_Module): 
            )



    def __setattr__(self, name, value):

        if isinstance(value, (borch.RandomVariable, borch.RVPair)):

        if isinstance(value, borch.RandomVariable):

            if "_parameters" not in self.__dict__:

                msg = (

                    f"Cannot assign random variables before"

                    f" {type(self).__name__}.__init__() call"

                )

                raise AttributeError(msg)

            q_rv = None

            if isinstance(value, borch.RVPair):

                q_rv = value.posterior

                value = value.prior

            if hasattr(self, "posterior") and self.posterior is not None:

                self.posterior.set_random_variable(name, value, q_rv)

                self.posterior.set_random_variable(name, value)

            setattr(self.prior, name, value)

            self._used_rvs.add(name)

        else:

src/borch/posterior/posterior.py

+1 −11
Original line number Diff line number Diff line
@@ -5,13 +5,6 @@ from torch.nn import Module 
class Posterior(Module):

    """Base class to create a posterior"""



    def register_q_random_variable(

        self, name, p_rv, q_rv

    ):  # pylint: disable=unused-argument

        """Allow the user to directly communicate with the posterior

        of what q_dist to use"""

        setattr(self, name, q_rv)



    def register_random_variable(self, name, rv):

        r"""Called the first time a RandomVariable gets added"""

        raise NotImplementedError
@@ -19,12 +12,9 @@ class Posterior(Module): 
    def update_random_variable(self, name, rv):

        """Code that gets run every time a random variable is added to a model"""



    def set_random_variable(self, name, rv, q_rv=None):

    def set_random_variable(self, name, rv):

        """If __setattr__ gets called on a model it will call here"""

        if not hasattr(self, name):

            if q_rv is not None:

                self.register_q_random_variable(name, rv, q_rv)

            else:

            self.register_random_variable(name, rv)

        self.update_random_variable(name, rv)

        return getattr(self, name)

src/borch/random_variable.py

+15 −8
Original line number Diff line number Diff line
@@ -9,6 +9,7 @@ from typing import Optional, Any, Union 
import torch



import borch.graph as graph

import borch



_VALIDATE_ARGS = contextvars.ContextVar("VALIDATE_ARGS", default=None)
@@ -215,15 +216,21 @@ class RandomVariable(graph.Graph): 
RVORDIST = Union[torch.distributions.Distribution, RandomVariable]





class RVPair:

class RVPair(graph.Graph):

    """

    Provide a prior and posterior pair to the model where the posterior is

    communicated to the posterior as a suggestion?

    Provide a prior and the corresponding approximating

    distribution.



    works like Rv(.., q_dist=) did before

    This is useful when one wants a custom approximating

    distribution.

    """



    def __init__(self, prior, posterior):

        # TODO do we want different arg names

        self.prior = prior

        self.posterior = posterior
 
    def __init__(self, p_dist, q_dist):

        posterior = borch.posterior.Manual()

        posterior.distribution = q_dist

        super().__init__(posterior=posterior)

        self.distribution = p_dist



    def forward(self):

        """The forward"""

        return self.distribution