def test_flatten_rate(): n = 1000 init_data = np.random.randint(1 << 16, size=8 * n, dtype='uint64') init_message = cs.base_message((1, )) for datum in init_data: init_message = cs.Uniform(16).push(init_message, datum) l_init = len(cs.flatten(init_message)) ps = np.random.rand(n, 1) data = np.random.rand(n, 1) < ps message = init_message for p, datum in zip(ps, data): message = cs.Bernoulli(p, 14).push(message, datum) l_scalar = len(cs.flatten(message)) message = init_message message = cs.reshape_head(message, (n, 1)) message = cs.Bernoulli(ps, 14).push(message, data) l_vector = len(cs.flatten(message)) assert (l_vector - l_init) / (l_scalar - l_init) - 1 < 0.001
def test_bernoulli(): precision = 4 shape = (2, 3, 5) p = rng.random(shape) data = np.uint64(rng.random(shape) < p) check_codec(shape, cs.Bernoulli(p, precision), data)
latent_dim = 40 latent_shape = (batch_size, latent_dim) latent_size = np.prod(latent_shape) obs_shape = (batch_size, 28 * 28) obs_size = np.prod(obs_shape) ## Setup codecs # VAE codec model = BinaryVAE(hidden_dim=100, latent_dim=40) model.load_state_dict(torch.load('vae_params')) rec_net = torch_fun_to_numpy_fun(model.encode) gen_net = torch_fun_to_numpy_fun(model.decode) obs_codec = lambda p: cs.Bernoulli(p, bernoulli_precision) def vae_view(head): return ag_tuple((np.reshape(head[:latent_size], latent_shape), np.reshape(head[latent_size:], obs_shape))) vae_append, vae_pop = cs.repeat(cs.substack( bb_ans.VAE(gen_net, rec_net, obs_codec, prior_precision, q_precision), vae_view), num_batches) ## Load mnist images images = datasets.MNIST(sys.argv[1], train=False, download=True).data.numpy() images = np.uint64(rng.random_sample(np.shape(images)) < images / 255.) images = np.split(np.reshape(images, (num_images, -1)), num_batches) ## Encode