def test_vae():
minibatch_size = 10
data = pm.floatX(np.random.rand(100))
x_mini = pm.Minibatch(data, minibatch_size)
x_inp = at.vector()
x_inp.tag.test_value = data[:minibatch_size]
ae = aesara.shared(pm.floatX([0.1, 0.1]))
be = aesara.shared(pm.floatX(1.0))
ad = aesara.shared(pm.floatX(1.0))
bd = aesara.shared(pm.floatX(1.0))
enc = x_inp.dimshuffle(0, "x") * ae.dimshuffle("x", 0) + be
mu, rho = enc[:, 0], enc[:, 1]
with pm.Model():
# Hidden variables
zs = pm.Normal("zs", mu=0, sigma=1, size=minibatch_size)
dec = zs * ad + bd
# Observation model
pm.Normal("xs_", mu=dec, sigma=0.1, observed=x_inp)
pm.fit(
1,
local_rv={zs: dict(mu=mu, rho=rho)},
more_replacements={x_inp: x_mini},
more_obj_params=[ae, be, ad, bd],
)
def test_cloning_available(self):
gop = pm.Minibatch(np.arange(100), 1)
res = gop**2
shared = aesara.shared(np.array([10]))
res1 = aesara.clone_replace(res, {gop: shared})
f = aesara.function([], res1)
assert f() == np.array([100])
def test_align(self):
m = pm.Minibatch(np.arange(1000), 1, random_seed=1)
n = pm.Minibatch(np.arange(1000), 1, random_seed=1)
f = aesara.function([], [m, n])
n.eval() # not aligned
a, b = zip(*(f() for _ in range(1000)))
assert a != b
pm.align_minibatches()
a, b = zip(*(f() for _ in range(1000)))
assert a == b
n.eval() # not aligned
pm.align_minibatches([m])
a, b = zip(*(f() for _ in range(1000)))
assert a != b
pm.align_minibatches([m, n])
a, b = zip(*(f() for _ in range(1000)))
assert a == b
def simple_model_data(use_minibatch):
n = 1000
sigma0 = 2.0
mu0 = 4.0
sigma = 3.0
mu = -5.0
data = sigma * np.random.randn(n) + mu
d = n / sigma**2 + 1 / sigma0**2
mu_post = (n * np.mean(data) / sigma**2 + mu0 / sigma0**2) / d
if use_minibatch:
data = pm.Minibatch(data)
return dict(
n=n,
data=data,
mu_post=mu_post,
d=d,
mu0=mu0,
sigma0=sigma0,
sigma=sigma,
)
def test_special_batch_size(self, batch_size, expected):
mb = pm.Minibatch(self.data, batch_size)
assert mb.eval().shape == expected
def test_2d(self):
mb = pm.Minibatch(self.data, [(10, 42), (4, 42)])
assert mb.eval().shape == (10, 4, 40, 10, 50)
def test_1d(self):
mb = pm.Minibatch(self.data, 20)
assert mb.eval().shape == (20, 10, 40, 10, 50)
def test_mixed2(self):
with pm.Model():
data = np.random.rand(10, 20, 30, 40, 50)
mb = pm.Minibatch(data, [2, None, 20])
Normal("n", observed=mb, total_size=(10, None, 30))
def test_special4(self):
mb = pm.Minibatch(self.data, [10, None, Ellipsis, (4, 42)])
assert mb.eval().shape == (10, 10, 40, 10, 4)
def test_special1(self):
mb = pm.Minibatch(self.data, [(10, 42), None, (4, 42)])
assert mb.eval().shape == (10, 10, 4, 10, 50)
