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)