<jupyter_start><jupyter_text>Same as above but batched over latent variables<jupyter_code>n_steps = 100
n_sub_steps = 10

hyper_chain = []
latent_chain = []

# set up batching of latent parameters
latent_batch_size = 64
n_latent_batches = n_latent // latent_batch_size + int( (n_latent % latent_batch_size) > 0)
batch_indices = [np.arange(latent_batch_size*i, min(latent_batch_size*(i+1), n_latent)) for i in range(n_latent_batches)]

# main loop
with trange(n_steps) as t:
    for step in t:
    
        # sample hyper-parameters, conditioned on latent parameters
        hyper_samples_ = affine_sample(log_hyperparameter_conditional, n_sub_steps, hyper_current_state, args=[latent_current_state[np.random.randint(0, 2)][np.random.randint(0, n_latent_walkers),...]], progressbar=False)
        hyper_current_state = tf.split(hyper_samples_[-1,...], (n_hyper_walkers, n_hyper_walkers), axis=0)

        # sample latent parameters, conditioned on hyper-parameters (do it in batches and concatenate them together)
        latent_samples_ = tf.concat([affine_sample_batch(log_latentparameter_conditional, 
                                                         n_sub_steps, 
                                                         [tf.gather(latent_current_state[0], batch_indices[_], axis=1), tf.gather(latent_current_state[1], batch_indices[_], axis=1)], 
                                                         args=[hyper_current_state[np.random.randint(0, 2)][np.random.randint(0, n_hyper_walkers),...], tf.gather(y_, batch_indices[_], axis=0)],
                                                         progressbar=False) for _ in range(n_latent_batches)], axis=2)
        latent_current_state = tf.split(latent_samples_[-1,...], (n_hyper_walkers, n_hyper_walkers), axis=0)
        
        # append samples to the chains
        hyper_chain.append(tf.concat(hyper_current_state, axis=0))
        latent_chain.append(tf.concat(latent_current_state, axis=0))
        
# stack the chains
hyper_chain = tf.stack(hyper_chain, axis=0)
latent_chain = tf.stack(latent_chain, axis=0)
plt.hist(hyper_chain[50:,:,0].numpy().flatten(), bins = 50, alpha = 0.5)
plt.axvline(mu_, ls = '--')
plt.show()

plt.hist(np.exp(hyper_chain[50:,:,1].numpy().flatten()), bins = 50, alpha = 0.5)
plt.axvline(np.exp(logsigma_), ls = '--')
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12420517168291649, 'fraction_numerical': 0.018228062738448497, 'mean_word_length': 3.1843971631205674, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '664415', 'license': 'MIT', 'n_tokens_mistral': 814, 'n_tokens_neox': 766, 'n_words': 138, 'path': 'affine_gibbs_sampling_example.ipynb', 'repo_name': 'justinalsing/affine', 'size': '2359'}