<jupyter_start><jupyter_text>Compare two chains<jupyter_code>all_pred_y_stack_ = onp.load("predictive_1.npz")["preds"]
pred_mean = onp.mean(all_pred_y_stack_, axis=0)
pred_std = onp.std(all_pred_y_stack_, axis=0)

plt.plot(x, y, "bo")
plt.plot(x_, y_, "-k")
plt.fill_between(x_, pred_mean + 3*pred_std, pred_mean - 3*pred_std, color="g", alpha=0.3)
for preds in all_pred_y_stack_[::20]:
    plt.plot(x_, preds, alpha=0.5, color="g")
all_pred_y_stack_ = onp.load("predictive_0.npz")["preds"]
pred_mean = onp.mean(all_pred_y_stack_, axis=0)
pred_std = onp.std(all_pred_y_stack_, axis=0)

plt.plot(x, y, "bo")
plt.plot(x_, y_, "-k")
plt.fill_between(x_, pred_mean + 3*pred_std, pred_mean - 3*pred_std, color="b", alpha=0.3)
for preds in all_pred_y_stack_[::20]:
    plt.plot(x_, preds, alpha=0.5, color="b")
all_pred_y_stack_ = onp.load("predictive_1.npz")["preds"]
pred_mean = onp.mean(all_pred_y_stack_, axis=0)
pred_std = onp.std(all_pred_y_stack_, axis=0)

plt.plot(x, y, "bo")
plt.plot(x_, y_, "-k")
plt.fill_between(x_, pred_mean + 3*pred_std, pred_mean - 3*pred_std, color="g", alpha=0.3)
for preds in all_pred_y_stack_[::20]:
    plt.plot(x_, preds, alpha=0.5, color="g")
    
    all_pred_y_stack_ = onp.load("predictive_0.npz")["preds"]
pred_mean = onp.mean(all_pred_y_stack_, axis=0)
pred_std = onp.std(all_pred_y_stack_, axis=0)

plt.fill_between(x_, pred_mean + 3*pred_std, pred_mean - 3*pred_std, color="b", alpha=0.3)
for preds in all_pred_y_stack_[::20]:
    plt.plot(x_, preds, alpha=0.5, color="b")
pred_y_, _ = net_apply(params, net_state, None, (f_, None), False)
plt.plot(x, y, "bo")
plt.plot(x_, y_, "-k")
plt.fill_between(x_, pred_mean + pred_std, pred_mean - pred_std, color="g", alpha=0.5)
plt.plot(x_, all_pred_y_[-1], "--g")
for preds in all_pred_y_stack_[::50]:
    plt.plot(x_, preds, alpha=0.5, color="g")
# onp.savez("predictive_1.npz", preds=all_pred_y_stack_)
# onp.savez("predictive_0.npz", preds=all_pred_y_stack_)
pred_y_, _ = net_apply(params, net_state, None, (f_, None), False)
plt.plot(x, y, "bo")
# plt.fill_between(x_, pred_mean + pred_std, pred_mean - pred_std, color="g", alpha=0.5)
for preds in all_pred_y_stack_[::1]:
    plt.plot(x_, preds, alpha=0.5, color="g")<jupyter_output><empty_output><jupyter_text>Loss Surface<jupyter_code>def tree_diff(a, b):
    return jax.tree_multimap(lambda p_a, p_b: p_a - p_b, a, b)
def tree_norm(a):
    return float(jnp.sqrt(sum([jnp.sum(p_a**2) for p_a in jax.tree_leaves(a)])))
    
def get_u_v_o(params1, params2, params3):
    
#     print([p.shape for p in jax.tree_leaves(params1)])
    
    u_params = tree_diff(params2, params1)
    u_norm = tree_norm(u_params)
    u_params = tree_utils.tree_scalarmul(u_params, 1 / u_norm)
    v_params = tree_diff(params3, params1)
    uv_dot = tree_utils.tree_dot(u_params, v_params)
    v_params = jax.tree_multimap(lambda v, u: v - uv_dot * u, v_params, u_params)
    v_norm = tree_norm(v_params)
    v_params = tree_utils.tree_scalarmul(v_params, 1 / v_norm)
    
    return u_params, u_norm, v_params, v_norm, params1
checkpoint_dict = checkpoint_utils.load_checkpoint("model_step_100.pt")
params1 = checkpoint_dict["params"]

checkpoint_dict = checkpoint_utils.load_checkpoint("model_step_150.pt")
params2 = checkpoint_dict["params"]

checkpoint_dict = checkpoint_utils.load_checkpoint("model_step_200.pt")
params3 = checkpoint_dict["params"]
for params in [params1, params2, params3]:
    log_prob, _, _, _ = evaluate(params, net_state, train_set)
    print(log_prob)
u_params, u_norm, v_params, v_norm, o_params = get_u_v_o(params1, params2, params3)

u_ts = onp.linspace(-.25, 1.25, 40)
v_ts = onp.linspace(-.25, 1.25, 40)
n_u, n_v = len(u_ts), len(v_ts)
log_probs = onp.zeros((n_u, n_v))
grid = onp.zeros((n_u, n_v, 2))

for u_i, u_t in enumerate(tqdm.tqdm(u_ts)):
    for v_i, v_t in enumerate(v_ts):
        grid[u_i, v_i] = [u_t, v_t]
        params = jax.tree_multimap(lambda u, v, o: o + u * u_t * u_norm + v * v_t * v_norm, 
                                   u_params, v_params, o_params)
        log_prob, _, _, _ = evaluate(params, net_state, train_set)
        log_probs[u_i, v_i] = float(log_prob)
log_probs_cliped = onp.clip(log_probs, -7.e4, 0)
plt.contour(grid[:, :, 0], grid[:, :, 1], log_probs_cliped, zorder=1)
plt.contourf(grid[:, :, 0], grid[:, :, 1], log_probs_cliped, zorder=0, alpha=0.55)
plt.plot([0., 1., 0.5], [0., 0., 1.], "ro", ms=20, markeredgecolor="k")<jupyter_output><empty_output><jupyter_text>Distance analysis<jupyter_code>checkpoint_dict = checkpoint_utils.load_checkpoint("model_step_100.pt")
params1 = checkpoint_dict["params"]

checkpoint_dict = checkpoint_utils.load_checkpoint("model_step_200.pt")
params2 = checkpoint_dict["params"]
d = sum([jnp.size(p) for p in jax.tree_leaves(params1)])
theory_dist = jnp.sqrt(2 * d * prior_std**2)

dist = tree_norm(tree_diff(params2, params1))
print(dist, theory_dist)
a,b,c = zip(*[[1, "a", 0.1], [2, "b", 0.2], [3, "c", 0.3]])<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.16579476861167003, 'fraction_numerical': 0.02937625754527163, 'mean_word_length': 5.3486590038314175, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '572433', 'license': 'Apache-2.0', 'n_tokens_mistral': 2245, 'n_tokens_neox': 2165, 'n_words': 381, 'path': 'bnn_hmc/notebooks/synthetic_regression_hmc.ipynb', 'repo_name': 'izmailovpavel/google-research', 'size': '4970'}