Simple wrapper custom variable

1

Hi. I’m new to Edward and I’m confused about how to set up custom random variables.

I figure the “hello world” of custom random variables is to just set up a basic wrapper to an existing Edward random variable class. Here is my attempt to do that with NormalWithSoftplusScale:

import edward as ed
import numpy as np
import tensorflow as tf

from edward.models import RandomVariable
from tensorflow.contrib.distributions import Distribution

from edward.models import Normal
from edward.models import NormalWithSoftplusScale

class CustomRandomVariable(RandomVariable, Distribution):
    def __init__(self, loc, scale, sample_shape,
                  validate_args=False,
                  allow_nan_stats=True,
                  name="Custom"):
        self._loc = loc
        self._scale = scale
        self._sample_shape = sample_shape
        super(CustomRandomVariable, self).__init__(
            dtype=tf.float32,
            reparameterization_type=tf.contrib.distributions.FULLY_REPARAMETERIZED,
            validate_args=validate_args,
            allow_nan_stats=allow_nan_stats,
            sample_shape=sample_shape,
            name=name)

    def _log_prob(self, value):
        return NormalWithSoftplusScale(loc=self._loc, scale=self._scale, sample_shape=self._sample_shape)._log_prob(value)

    def _sample_n(self, n, seed=None):
        return NormalWithSoftplusScale(loc=self._loc, scale=self._scale, sample_shape=self._sample_shape)._sample_n(n)

X_train = np.random.normal(0.0, 1.0, 10)

mu = Normal(loc=10.0, scale=5.0)
sigma = Normal(loc=0.0, scale=1.0)

qmu = Normal(loc=tf.Variable(10.0), scale=tf.Variable(5.0))
qsigma = Normal(loc=tf.Variable(0.0), scale=tf.Variable(1.0))

X = CustomRandomVariable(loc=mu, scale=sigma, sample_shape=X_train.shape)
inference = ed.KLqp({mu:qmu, sigma:qsigma}, {X:X_train})
inference.run(n_samples=10, n_iter=250)

Running this gives the error:
TypeError: __init__() got an unexpected keyword argument 'reparameterization_type'

If I comment out the reparameterization_type argument in the super() call, I get a conflicting error:
TypeError: __init__() missing 1 required positional argument: 'reparameterization_type'

Damned if I do, damned if I don’t. I suppose there is a conflict in the arguments expected by the RandomVariable and Distribution classes. But I’ve seen examples that seem to make similar things work with the same type of call to super() in __init__(). So, I’m not sure of the best/proper way to resolve the arguments to make this wrapper class work.

I’m trying to implement a more complicated custom random variable, but I keep running into errors like the above. If someone can help be debug the example above, I think that will be a good start.

Thanks!

2

Thanks for raising this. The docs need to be updated; the inheritance is actually a done a little differently in, say, empirical.py and point_mass.py (see their source). This issue is also raised in https://github.com/blei-lab/edward/issues/755.

3

Thank you for the reply and the link, Dustin.

I have looked at the two examples (empirical.py and point_mass.py) and am still confused about how to construct a custom variable. In the examples, a class is constructed that’s a child class of Distribution. Then, this child class is used as a parent class for the CustomRandomVariable class. I’m not sure of the correct way to set this up so that I can pass parameters. Here is another attempt:

import edward as ed
import numpy as np
import tensorflow as tf

from edward.models import RandomVariable
from tensorflow.contrib.distributions import Distribution

from edward.models import Normal
from edward.models import NormalWithSoftplusScale

class distributions_CustomRandomVariable(Distribution):
    def __init__(self, loc, scale, sample_shape,
                  validate_args=False,
                  allow_nan_stats=True,
                  name="Custom"):
        self._loc = loc
        self._scale = scale
        self._sample_shape = sample_shape
        super(distributions_CustomRandomVariable, self).__init__(
            dtype=tf.float32,
            reparameterization_type=tf.contrib.distributions.FULLY_REPARAMETERIZED,
            validate_args=validate_args,
            allow_nan_stats=allow_nan_stats,
            name=name)

    def _log_prob(self, value):
        return NormalWithSoftplusScale(loc=self._loc, scale=self._scale, sample_shape=self._sample_shape)._log_prob(value)

    def _sample_n(self, n, seed=None):
        return NormalWithSoftplusScale(loc=self._loc, scale=self._scale, sample_shape=self._sample_shape)._sample_n(n, seed=seed)
    
class CustomRandomVariable(RandomVariable, distributions_CustomRandomVariable):
    def __init__(self, loc, scale, sample_shape):
        #super(CustomRandomVariable,self).__init__(loc, scale, sample_shape)
        #distributions_CustomRandomVariable.__init__(self, loc, scale, sample_shape)
        RandomVariable.__init__(self, loc, scale, sample_shape)
        

X_train = np.random.normal(0.0, 1.0, 10)

mu = Normal(loc=10.0, scale=5.0)
sigma = Normal(loc=0.0, scale=1.0)

qmu = Normal(loc=tf.Variable(10.0), scale=tf.Variable(5.0))
qsigma = Normal(loc=tf.Variable(0.0), scale=tf.Variable(1.0))

X = CustomRandomVariable(loc=mu, scale=sigma, sample_shape=X_train.shape)
inference = ed.KLqp({mu:qmu, sigma:qsigma}, {X:X_train})
inference.run(n_samples=10, n_iter=250)

This version fails with the error
TypeError: Key-value pair in data does not have same shape: (), (10,)

I think this error means that sample_shape is not being set properly in the CustomRandomVariable. I’ve tried using different __init__ strategies (see the commented out lines above) and I’ve tried setting the attribute directly in the distributions_CustomRandomVariable.__init__() function which results in the error
AttributeError: can't set attribute

I’m having trouble generalizing from the examples to get around these errors and I’m not really understanding how the inheritance is supposed to work. So if anyone can help me get this simple example working, I would be grateful.

Thanks!

4

sample_shape in this custom random variable overlaps with the base class RandomVariable’s sample_shape arg. It’s also not clear why you need sample_shape in the custom random variable if you only call class methods like log prob and sample from it.

Specifically, the following runs correctly:

import edward as ed
import numpy as np
import tensorflow as tf

from edward.models import RandomVariable
from tensorflow.contrib.distributions import Distribution

from edward.models import Normal
from edward.models import NormalWithSoftplusScale

class distributions_CustomRandomVariable(Distribution):
    def __init__(self, loc, scale,
                  validate_args=False,
                  allow_nan_stats=True,
                  name="Custom"):
        self._loc = loc
        self._scale = scale
        super(distributions_CustomRandomVariable, self).__init__(
            dtype=tf.float32,
            reparameterization_type=tf.contrib.distributions.FULLY_REPARAMETERIZED,
            validate_args=validate_args,
            allow_nan_stats=allow_nan_stats,
            name=name)

    def _log_prob(self, value):
        return NormalWithSoftplusScale(loc=self._loc, scale=self._scale)._log_prob(value)

    def _sample_n(self, n, seed=None):
        return NormalWithSoftplusScale(loc=self._loc, scale=self._scale)._sample_n(n, seed=seed)

class CustomRandomVariable(RandomVariable, distributions_CustomRandomVariable):
    def __init__(self, *args, **kwargs):
        #super(CustomRandomVariable,self).__init__(loc, scale, sample_shape)
        #distributions_CustomRandomVariable.__init__(self, loc, scale, sample_shape)
        RandomVariable.__init__(self, *args, **kwargs)


X_train = np.random.normal(0.0, 1.0, 10)

mu = Normal(loc=10.0, scale=5.0)
sigma = Normal(loc=0.0, scale=1.0)

qmu = Normal(loc=tf.Variable(10.0), scale=tf.Variable(5.0))
qsigma = Normal(loc=tf.Variable(0.0), scale=tf.Variable(1.0))

X = CustomRandomVariable(loc=mu, scale=sigma, sample_shape=X_train.shape)
print(X)
inference = ed.KLqp({mu:qmu, sigma:qsigma}, {X:X_train})
inference.run(n_samples=10, n_iter=250)
1 Like
5

Thanks, Dustin! Your example works for me. I think I understand some of my points of confusion with the custom distributions better now.

1 Like