Bayesian computations with Neural Networks

Neural network training is typically done with maximum likelihood estimation. Given the number of parameter invariances in neural network architectures, this often introduces a large number of local minima, making global optimisation very difficult.

JAXNS can easily compute the evidence and posterior of a neural network. The larger number of parameters means that we should be careful with how accurate we want this to be. With high precision, we’ll need more likelihood evaluations.

What we’ll do in this notebook

  1. Define a neural network model

  2. Find its maximum likelihood parameters using Global Optimisation

  3. Redefine it as a Bayesian neural network allowing biases to be sampled from a prior.

  4. Compute the evidence and posterior of the parameters.

Data

We’ll use the N-bit majority problem as our data. This is a binary classification problem where the input is a sequence of bits and the output is 1 if the majority of the bits are 1, and 0 otherwise.

[1]:
import os

from jax import tree_util

os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=6"

try:
    import haiku as hk
except ImportError:
    print("You must `pip install dm-haiku` first.")
    raise

from itertools import product

import jax
import tensorflow_probability.substrates.jax as tfp
from jax import random, numpy as jnp
from jax import vmap

from jaxns import resample

tfpd = tfp.distributions


def pytree_unravel(example_tree):
    """
    Returns functions to ravel and unravel a pytree.
    """
    leaf_list, tree_def = tree_util.tree_flatten(example_tree)

    sizes = [leaf.size for leaf in leaf_list]
    shapes = [leaf.shape for leaf in leaf_list]

    def ravel_fun(pytree):
        leaf_list, tree_def = tree_util.tree_flatten(pytree)
        return jnp.concatenate([leaf.ravel() for leaf in leaf_list])

    def unravel_fun(flat_array):
        leaf_list = []
        start = 0
        for size, shape in zip(sizes, shapes):
            leaf_list.append(flat_array[start:start + size].reshape(shape))
            start += size
        return tree_util.tree_unflatten(tree_def, leaf_list)

    return ravel_fun, unravel_fun

INFO[2024-01-10 16:55:05,659]: Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
INFO[2024-01-10 16:55:05,660]: Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
INFO[2024-01-10 16:55:05,661]: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: libtpu.so: cannot open shared object file: No such file or directory
WARNING[2024-01-10 16:55:05,662]: An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
[2]:
# Generate data

def n_bit_majority(x):
    return jnp.sum(x) > (x.size / 2)


num_variables = 7
options = [True, False]
x = jnp.asarray(list(product(options, repeat=num_variables)))  #N,2
y = vmap(n_bit_majority)(x)[:, None]  #N, 1
x = x.astype(jnp.float32)
print("Data:")

for input, output in zip(x, y):
    print(f"{input} -> {output}")

Data:
[1. 1. 1. 1. 1. 1. 1.] -> [ True]
[1. 1. 1. 1. 1. 1. 0.] -> [ True]
[1. 1. 1. 1. 1. 0. 1.] -> [ True]
[1. 1. 1. 1. 1. 0. 0.] -> [ True]
[1. 1. 1. 1. 0. 1. 1.] -> [ True]
[1. 1. 1. 1. 0. 1. 0.] -> [ True]
[1. 1. 1. 1. 0. 0. 1.] -> [ True]
[1. 1. 1. 1. 0. 0. 0.] -> [ True]
[1. 1. 1. 0. 1. 1. 1.] -> [ True]
[1. 1. 1. 0. 1. 1. 0.] -> [ True]
[1. 1. 1. 0. 1. 0. 1.] -> [ True]
[1. 1. 1. 0. 1. 0. 0.] -> [ True]
[1. 1. 1. 0. 0. 1. 1.] -> [ True]
[1. 1. 1. 0. 0. 1. 0.] -> [ True]
[1. 1. 1. 0. 0. 0. 1.] -> [ True]
[1. 1. 1. 0. 0. 0. 0.] -> [False]
[1. 1. 0. 1. 1. 1. 1.] -> [ True]
[1. 1. 0. 1. 1. 1. 0.] -> [ True]
[1. 1. 0. 1. 1. 0. 1.] -> [ True]
[1. 1. 0. 1. 1. 0. 0.] -> [ True]
[1. 1. 0. 1. 0. 1. 1.] -> [ True]
[1. 1. 0. 1. 0. 1. 0.] -> [ True]
[1. 1. 0. 1. 0. 0. 1.] -> [ True]
[1. 1. 0. 1. 0. 0. 0.] -> [False]
[1. 1. 0. 0. 1. 1. 1.] -> [ True]
[1. 1. 0. 0. 1. 1. 0.] -> [ True]
[1. 1. 0. 0. 1. 0. 1.] -> [ True]
[1. 1. 0. 0. 1. 0. 0.] -> [False]
[1. 1. 0. 0. 0. 1. 1.] -> [ True]
[1. 1. 0. 0. 0. 1. 0.] -> [False]
[1. 1. 0. 0. 0. 0. 1.] -> [False]
[1. 1. 0. 0. 0. 0. 0.] -> [False]
[1. 0. 1. 1. 1. 1. 1.] -> [ True]
[1. 0. 1. 1. 1. 1. 0.] -> [ True]
[1. 0. 1. 1. 1. 0. 1.] -> [ True]
[1. 0. 1. 1. 1. 0. 0.] -> [ True]
[1. 0. 1. 1. 0. 1. 1.] -> [ True]
[1. 0. 1. 1. 0. 1. 0.] -> [ True]
[1. 0. 1. 1. 0. 0. 1.] -> [ True]
[1. 0. 1. 1. 0. 0. 0.] -> [False]
[1. 0. 1. 0. 1. 1. 1.] -> [ True]
[1. 0. 1. 0. 1. 1. 0.] -> [ True]
[1. 0. 1. 0. 1. 0. 1.] -> [ True]
[1. 0. 1. 0. 1. 0. 0.] -> [False]
[1. 0. 1. 0. 0. 1. 1.] -> [ True]
[1. 0. 1. 0. 0. 1. 0.] -> [False]
[1. 0. 1. 0. 0. 0. 1.] -> [False]
[1. 0. 1. 0. 0. 0. 0.] -> [False]
[1. 0. 0. 1. 1. 1. 1.] -> [ True]
[1. 0. 0. 1. 1. 1. 0.] -> [ True]
[1. 0. 0. 1. 1. 0. 1.] -> [ True]
[1. 0. 0. 1. 1. 0. 0.] -> [False]
[1. 0. 0. 1. 0. 1. 1.] -> [ True]
[1. 0. 0. 1. 0. 1. 0.] -> [False]
[1. 0. 0. 1. 0. 0. 1.] -> [False]
[1. 0. 0. 1. 0. 0. 0.] -> [False]
[1. 0. 0. 0. 1. 1. 1.] -> [ True]
[1. 0. 0. 0. 1. 1. 0.] -> [False]
[1. 0. 0. 0. 1. 0. 1.] -> [False]
[1. 0. 0. 0. 1. 0. 0.] -> [False]
[1. 0. 0. 0. 0. 1. 1.] -> [False]
[1. 0. 0. 0. 0. 1. 0.] -> [False]
[1. 0. 0. 0. 0. 0. 1.] -> [False]
[1. 0. 0. 0. 0. 0. 0.] -> [False]
[0. 1. 1. 1. 1. 1. 1.] -> [ True]
[0. 1. 1. 1. 1. 1. 0.] -> [ True]
[0. 1. 1. 1. 1. 0. 1.] -> [ True]
[0. 1. 1. 1. 1. 0. 0.] -> [ True]
[0. 1. 1. 1. 0. 1. 1.] -> [ True]
[0. 1. 1. 1. 0. 1. 0.] -> [ True]
[0. 1. 1. 1. 0. 0. 1.] -> [ True]
[0. 1. 1. 1. 0. 0. 0.] -> [False]
[0. 1. 1. 0. 1. 1. 1.] -> [ True]
[0. 1. 1. 0. 1. 1. 0.] -> [ True]
[0. 1. 1. 0. 1. 0. 1.] -> [ True]
[0. 1. 1. 0. 1. 0. 0.] -> [False]
[0. 1. 1. 0. 0. 1. 1.] -> [ True]
[0. 1. 1. 0. 0. 1. 0.] -> [False]
[0. 1. 1. 0. 0. 0. 1.] -> [False]
[0. 1. 1. 0. 0. 0. 0.] -> [False]
[0. 1. 0. 1. 1. 1. 1.] -> [ True]
[0. 1. 0. 1. 1. 1. 0.] -> [ True]
[0. 1. 0. 1. 1. 0. 1.] -> [ True]
[0. 1. 0. 1. 1. 0. 0.] -> [False]
[0. 1. 0. 1. 0. 1. 1.] -> [ True]
[0. 1. 0. 1. 0. 1. 0.] -> [False]
[0. 1. 0. 1. 0. 0. 1.] -> [False]
[0. 1. 0. 1. 0. 0. 0.] -> [False]
[0. 1. 0. 0. 1. 1. 1.] -> [ True]
[0. 1. 0. 0. 1. 1. 0.] -> [False]
[0. 1. 0. 0. 1. 0. 1.] -> [False]
[0. 1. 0. 0. 1. 0. 0.] -> [False]
[0. 1. 0. 0. 0. 1. 1.] -> [False]
[0. 1. 0. 0. 0. 1. 0.] -> [False]
[0. 1. 0. 0. 0. 0. 1.] -> [False]
[0. 1. 0. 0. 0. 0. 0.] -> [False]
[0. 0. 1. 1. 1. 1. 1.] -> [ True]
[0. 0. 1. 1. 1. 1. 0.] -> [ True]
[0. 0. 1. 1. 1. 0. 1.] -> [ True]
[0. 0. 1. 1. 1. 0. 0.] -> [False]
[0. 0. 1. 1. 0. 1. 1.] -> [ True]
[0. 0. 1. 1. 0. 1. 0.] -> [False]
[0. 0. 1. 1. 0. 0. 1.] -> [False]
[0. 0. 1. 1. 0. 0. 0.] -> [False]
[0. 0. 1. 0. 1. 1. 1.] -> [ True]
[0. 0. 1. 0. 1. 1. 0.] -> [False]
[0. 0. 1. 0. 1. 0. 1.] -> [False]
[0. 0. 1. 0. 1. 0. 0.] -> [False]
[0. 0. 1. 0. 0. 1. 1.] -> [False]
[0. 0. 1. 0. 0. 1. 0.] -> [False]
[0. 0. 1. 0. 0. 0. 1.] -> [False]
[0. 0. 1. 0. 0. 0. 0.] -> [False]
[0. 0. 0. 1. 1. 1. 1.] -> [ True]
[0. 0. 0. 1. 1. 1. 0.] -> [False]
[0. 0. 0. 1. 1. 0. 1.] -> [False]
[0. 0. 0. 1. 1. 0. 0.] -> [False]
[0. 0. 0. 1. 0. 1. 1.] -> [False]
[0. 0. 0. 1. 0. 1. 0.] -> [False]
[0. 0. 0. 1. 0. 0. 1.] -> [False]
[0. 0. 0. 1. 0. 0. 0.] -> [False]
[0. 0. 0. 0. 1. 1. 1.] -> [False]
[0. 0. 0. 0. 1. 1. 0.] -> [False]
[0. 0. 0. 0. 1. 0. 1.] -> [False]
[0. 0. 0. 0. 1. 0. 0.] -> [False]
[0. 0. 0. 0. 0. 1. 1.] -> [False]
[0. 0. 0. 0. 0. 1. 0.] -> [False]
[0. 0. 0. 0. 0. 0. 1.] -> [False]
[0. 0. 0. 0. 0. 0. 0.] -> [False]
[3]:
from jaxns import Prior, Model


def prior_model():
    def compute_logits(x):
        mlp = hk.Sequential([hk.Linear(1),
                             jax.nn.sigmoid,
                             hk.Linear(1)])
        return mlp(x)

    fn = hk.without_apply_rng(hk.transform(compute_logits))
    init_params = fn.init(random.PRNGKey(0), x)

    ravel_fn, unravel_fn = pytree_unravel(init_params)
    ndims = ravel_fn(init_params).size
    flat_params = yield Prior(tfpd.Uniform(-10. * jnp.ones(ndims), 10. * jnp.ones(ndims)), name='flat_params')
    params = unravel_fn(flat_params)
    logits = fn.apply(params, x)
    return logits


def log_likelihood(logits):
    # Classification probelm, so we use a Bernoulli likelihood
    return tfpd.Bernoulli(logits=logits).log_prob(y).mean()


model = Model(prior_model=prior_model, log_likelihood=log_likelihood)



[4]:

model.sanity_check(random.PRNGKey(0), S=100)
INFO[2024-01-10 16:55:07,850]: Sanity check...
INFO[2024-01-10 16:55:08,319]: Sanity check passed
[5]:
from jaxns.experimental import DefaultGlobalOptimisation, GlobalOptimisationTerminationCondition

go = DefaultGlobalOptimisation(model=model,
                               num_search_chains=model.U_ndims * 10,
                               num_parallel_workers=len(jax.devices()))

results = go(
    random.PRNGKey(0),
    GlobalOptimisationTerminationCondition(log_likelihood_contour=-0.01,
                                           max_likelihood_evaluations=2e6)
)
go.summary(results)

INFO[2024-01-10 16:55:08,340]: Using 6 parallel workers, each running identical samplers.
--------
Termination Conditions:
--------
likelihood evals: 343262
samples: 14700
likelihood evals / sample: 23.4
--------
max(log_L)=-0.028
relative spread: 0.00063
absolute spread: 1.7e-05
--------
flat_params[#]: max(L) est.
flat_params[0]: -10.0
flat_params[1]: 2.91
flat_params[2]: 2.91
flat_params[3]: 2.91
flat_params[4]: 2.91
flat_params[5]: 2.9
flat_params[6]: 2.91
flat_params[7]: 2.91
flat_params[8]: -5.25
flat_params[9]: 10.0
--------
[6]:
logits = model.prepare_input(results.U_solution)[0]
predictions = jax.nn.sigmoid(logits)
for i in range(len(y)):
    pred = predictions[i] > 0.5
    print(f"{i}: {x[i]} -> {y[i]} | pred: {pred} {'✓' if pred == y[i] else '✗'}")

accuracy = jnp.mean((predictions > 0.5) == y)
print(f"Accuracy: {accuracy * 100:.1f}%")
0: [1. 1. 1. 1. 1. 1. 1.] -> [ True] | pred: [ True] ✓
1: [1. 1. 1. 1. 1. 1. 0.] -> [ True] | pred: [ True] ✓
2: [1. 1. 1. 1. 1. 0. 1.] -> [ True] | pred: [ True] ✓
3: [1. 1. 1. 1. 1. 0. 0.] -> [ True] | pred: [ True] ✓
4: [1. 1. 1. 1. 0. 1. 1.] -> [ True] | pred: [ True] ✓
5: [1. 1. 1. 1. 0. 1. 0.] -> [ True] | pred: [ True] ✓
6: [1. 1. 1. 1. 0. 0. 1.] -> [ True] | pred: [ True] ✓
7: [1. 1. 1. 1. 0. 0. 0.] -> [ True] | pred: [ True] ✓
8: [1. 1. 1. 0. 1. 1. 1.] -> [ True] | pred: [ True] ✓
9: [1. 1. 1. 0. 1. 1. 0.] -> [ True] | pred: [ True] ✓
10: [1. 1. 1. 0. 1. 0. 1.] -> [ True] | pred: [ True] ✓
11: [1. 1. 1. 0. 1. 0. 0.] -> [ True] | pred: [ True] ✓
12: [1. 1. 1. 0. 0. 1. 1.] -> [ True] | pred: [ True] ✓
13: [1. 1. 1. 0. 0. 1. 0.] -> [ True] | pred: [ True] ✓
14: [1. 1. 1. 0. 0. 0. 1.] -> [ True] | pred: [ True] ✓
15: [1. 1. 1. 0. 0. 0. 0.] -> [False] | pred: [False] ✓
16: [1. 1. 0. 1. 1. 1. 1.] -> [ True] | pred: [ True] ✓
17: [1. 1. 0. 1. 1. 1. 0.] -> [ True] | pred: [ True] ✓
18: [1. 1. 0. 1. 1. 0. 1.] -> [ True] | pred: [ True] ✓
19: [1. 1. 0. 1. 1. 0. 0.] -> [ True] | pred: [ True] ✓
20: [1. 1. 0. 1. 0. 1. 1.] -> [ True] | pred: [ True] ✓
21: [1. 1. 0. 1. 0. 1. 0.] -> [ True] | pred: [ True] ✓
22: [1. 1. 0. 1. 0. 0. 1.] -> [ True] | pred: [ True] ✓
23: [1. 1. 0. 1. 0. 0. 0.] -> [False] | pred: [False] ✓
24: [1. 1. 0. 0. 1. 1. 1.] -> [ True] | pred: [ True] ✓
25: [1. 1. 0. 0. 1. 1. 0.] -> [ True] | pred: [ True] ✓
26: [1. 1. 0. 0. 1. 0. 1.] -> [ True] | pred: [ True] ✓
27: [1. 1. 0. 0. 1. 0. 0.] -> [False] | pred: [False] ✓
28: [1. 1. 0. 0. 0. 1. 1.] -> [ True] | pred: [ True] ✓
29: [1. 1. 0. 0. 0. 1. 0.] -> [False] | pred: [False] ✓
30: [1. 1. 0. 0. 0. 0. 1.] -> [False] | pred: [False] ✓
31: [1. 1. 0. 0. 0. 0. 0.] -> [False] | pred: [False] ✓
32: [1. 0. 1. 1. 1. 1. 1.] -> [ True] | pred: [ True] ✓
33: [1. 0. 1. 1. 1. 1. 0.] -> [ True] | pred: [ True] ✓
34: [1. 0. 1. 1. 1. 0. 1.] -> [ True] | pred: [ True] ✓
35: [1. 0. 1. 1. 1. 0. 0.] -> [ True] | pred: [ True] ✓
36: [1. 0. 1. 1. 0. 1. 1.] -> [ True] | pred: [ True] ✓
37: [1. 0. 1. 1. 0. 1. 0.] -> [ True] | pred: [ True] ✓
38: [1. 0. 1. 1. 0. 0. 1.] -> [ True] | pred: [ True] ✓
39: [1. 0. 1. 1. 0. 0. 0.] -> [False] | pred: [False] ✓
40: [1. 0. 1. 0. 1. 1. 1.] -> [ True] | pred: [ True] ✓
41: [1. 0. 1. 0. 1. 1. 0.] -> [ True] | pred: [ True] ✓
42: [1. 0. 1. 0. 1. 0. 1.] -> [ True] | pred: [ True] ✓
43: [1. 0. 1. 0. 1. 0. 0.] -> [False] | pred: [False] ✓
44: [1. 0. 1. 0. 0. 1. 1.] -> [ True] | pred: [ True] ✓
45: [1. 0. 1. 0. 0. 1. 0.] -> [False] | pred: [False] ✓
46: [1. 0. 1. 0. 0. 0. 1.] -> [False] | pred: [False] ✓
47: [1. 0. 1. 0. 0. 0. 0.] -> [False] | pred: [False] ✓
48: [1. 0. 0. 1. 1. 1. 1.] -> [ True] | pred: [ True] ✓
49: [1. 0. 0. 1. 1. 1. 0.] -> [ True] | pred: [ True] ✓
50: [1. 0. 0. 1. 1. 0. 1.] -> [ True] | pred: [ True] ✓
51: [1. 0. 0. 1. 1. 0. 0.] -> [False] | pred: [False] ✓
52: [1. 0. 0. 1. 0. 1. 1.] -> [ True] | pred: [ True] ✓
53: [1. 0. 0. 1. 0. 1. 0.] -> [False] | pred: [False] ✓
54: [1. 0. 0. 1. 0. 0. 1.] -> [False] | pred: [False] ✓
55: [1. 0. 0. 1. 0. 0. 0.] -> [False] | pred: [False] ✓
56: [1. 0. 0. 0. 1. 1. 1.] -> [ True] | pred: [ True] ✓
57: [1. 0. 0. 0. 1. 1. 0.] -> [False] | pred: [False] ✓
58: [1. 0. 0. 0. 1. 0. 1.] -> [False] | pred: [False] ✓
59: [1. 0. 0. 0. 1. 0. 0.] -> [False] | pred: [False] ✓
60: [1. 0. 0. 0. 0. 1. 1.] -> [False] | pred: [False] ✓
61: [1. 0. 0. 0. 0. 1. 0.] -> [False] | pred: [False] ✓
62: [1. 0. 0. 0. 0. 0. 1.] -> [False] | pred: [False] ✓
63: [1. 0. 0. 0. 0. 0. 0.] -> [False] | pred: [False] ✓
64: [0. 1. 1. 1. 1. 1. 1.] -> [ True] | pred: [ True] ✓
65: [0. 1. 1. 1. 1. 1. 0.] -> [ True] | pred: [ True] ✓
66: [0. 1. 1. 1. 1. 0. 1.] -> [ True] | pred: [ True] ✓
67: [0. 1. 1. 1. 1. 0. 0.] -> [ True] | pred: [ True] ✓
68: [0. 1. 1. 1. 0. 1. 1.] -> [ True] | pred: [ True] ✓
69: [0. 1. 1. 1. 0. 1. 0.] -> [ True] | pred: [ True] ✓
70: [0. 1. 1. 1. 0. 0. 1.] -> [ True] | pred: [ True] ✓
71: [0. 1. 1. 1. 0. 0. 0.] -> [False] | pred: [False] ✓
72: [0. 1. 1. 0. 1. 1. 1.] -> [ True] | pred: [ True] ✓
73: [0. 1. 1. 0. 1. 1. 0.] -> [ True] | pred: [ True] ✓
74: [0. 1. 1. 0. 1. 0. 1.] -> [ True] | pred: [ True] ✓
75: [0. 1. 1. 0. 1. 0. 0.] -> [False] | pred: [False] ✓
76: [0. 1. 1. 0. 0. 1. 1.] -> [ True] | pred: [ True] ✓
77: [0. 1. 1. 0. 0. 1. 0.] -> [False] | pred: [False] ✓
78: [0. 1. 1. 0. 0. 0. 1.] -> [False] | pred: [False] ✓
79: [0. 1. 1. 0. 0. 0. 0.] -> [False] | pred: [False] ✓
80: [0. 1. 0. 1. 1. 1. 1.] -> [ True] | pred: [ True] ✓
81: [0. 1. 0. 1. 1. 1. 0.] -> [ True] | pred: [ True] ✓
82: [0. 1. 0. 1. 1. 0. 1.] -> [ True] | pred: [ True] ✓
83: [0. 1. 0. 1. 1. 0. 0.] -> [False] | pred: [False] ✓
84: [0. 1. 0. 1. 0. 1. 1.] -> [ True] | pred: [ True] ✓
85: [0. 1. 0. 1. 0. 1. 0.] -> [False] | pred: [False] ✓
86: [0. 1. 0. 1. 0. 0. 1.] -> [False] | pred: [False] ✓
87: [0. 1. 0. 1. 0. 0. 0.] -> [False] | pred: [False] ✓
88: [0. 1. 0. 0. 1. 1. 1.] -> [ True] | pred: [ True] ✓
89: [0. 1. 0. 0. 1. 1. 0.] -> [False] | pred: [False] ✓
90: [0. 1. 0. 0. 1. 0. 1.] -> [False] | pred: [False] ✓
91: [0. 1. 0. 0. 1. 0. 0.] -> [False] | pred: [False] ✓
92: [0. 1. 0. 0. 0. 1. 1.] -> [False] | pred: [False] ✓
93: [0. 1. 0. 0. 0. 1. 0.] -> [False] | pred: [False] ✓
94: [0. 1. 0. 0. 0. 0. 1.] -> [False] | pred: [False] ✓
95: [0. 1. 0. 0. 0. 0. 0.] -> [False] | pred: [False] ✓
96: [0. 0. 1. 1. 1. 1. 1.] -> [ True] | pred: [ True] ✓
97: [0. 0. 1. 1. 1. 1. 0.] -> [ True] | pred: [ True] ✓
98: [0. 0. 1. 1. 1. 0. 1.] -> [ True] | pred: [ True] ✓
99: [0. 0. 1. 1. 1. 0. 0.] -> [False] | pred: [False] ✓
100: [0. 0. 1. 1. 0. 1. 1.] -> [ True] | pred: [ True] ✓
101: [0. 0. 1. 1. 0. 1. 0.] -> [False] | pred: [False] ✓
102: [0. 0. 1. 1. 0. 0. 1.] -> [False] | pred: [False] ✓
103: [0. 0. 1. 1. 0. 0. 0.] -> [False] | pred: [False] ✓
104: [0. 0. 1. 0. 1. 1. 1.] -> [ True] | pred: [ True] ✓
105: [0. 0. 1. 0. 1. 1. 0.] -> [False] | pred: [False] ✓
106: [0. 0. 1. 0. 1. 0. 1.] -> [False] | pred: [False] ✓
107: [0. 0. 1. 0. 1. 0. 0.] -> [False] | pred: [False] ✓
108: [0. 0. 1. 0. 0. 1. 1.] -> [False] | pred: [False] ✓
109: [0. 0. 1. 0. 0. 1. 0.] -> [False] | pred: [False] ✓
110: [0. 0. 1. 0. 0. 0. 1.] -> [False] | pred: [False] ✓
111: [0. 0. 1. 0. 0. 0. 0.] -> [False] | pred: [False] ✓
112: [0. 0. 0. 1. 1. 1. 1.] -> [ True] | pred: [ True] ✓
113: [0. 0. 0. 1. 1. 1. 0.] -> [False] | pred: [False] ✓
114: [0. 0. 0. 1. 1. 0. 1.] -> [False] | pred: [False] ✓
115: [0. 0. 0. 1. 1. 0. 0.] -> [False] | pred: [False] ✓
116: [0. 0. 0. 1. 0. 1. 1.] -> [False] | pred: [False] ✓
117: [0. 0. 0. 1. 0. 1. 0.] -> [False] | pred: [False] ✓
118: [0. 0. 0. 1. 0. 0. 1.] -> [False] | pred: [False] ✓
119: [0. 0. 0. 1. 0. 0. 0.] -> [False] | pred: [False] ✓
120: [0. 0. 0. 0. 1. 1. 1.] -> [False] | pred: [False] ✓
121: [0. 0. 0. 0. 1. 1. 0.] -> [False] | pred: [False] ✓
122: [0. 0. 0. 0. 1. 0. 1.] -> [False] | pred: [False] ✓
123: [0. 0. 0. 0. 1. 0. 0.] -> [False] | pred: [False] ✓
124: [0. 0. 0. 0. 0. 1. 1.] -> [False] | pred: [False] ✓
125: [0. 0. 0. 0. 0. 1. 0.] -> [False] | pred: [False] ✓
126: [0. 0. 0. 0. 0. 0. 1.] -> [False] | pred: [False] ✓
127: [0. 0. 0. 0. 0. 0. 0.] -> [False] | pred: [False] ✓
Accuracy: 100.0%
[7]:


def prior_model(): bias_scale = yield Prior(tfpd.Uniform(0.1, 10.), name='bias_scale').parametrised() bias = yield Prior(tfpd.Normal(jnp.zeros(1), bias_scale * jnp.ones(1)), name='bias') def compute_logits(x): mlp = hk.Sequential([hk.Linear(1, with_bias=False), lambda x: x + bias, jax.nn.sigmoid, hk.Linear(1)]) return mlp(x) fn = hk.without_apply_rng(hk.transform(compute_logits)) init_params = fn.init(random.PRNGKey(0), x) ravel_fn, unravel_fn = pytree_unravel(init_params) ndims = ravel_fn(init_params).size # We'll parameterise the NN's params flat_params = yield Prior( tfpd.Uniform(-10. * jnp.ones(ndims), 10. * jnp.ones(ndims)), name='flat_params').parametrised() params = unravel_fn(flat_params) logits = fn.apply(params, x) return logits def log_likelihood(logits): # Classification probelm, so we use a Bernoulli likelihood return tfpd.Bernoulli(logits=logits).log_prob(y).mean() model = Model(prior_model=prior_model, log_likelihood=log_likelihood)
[8]:
from jaxns.experimental import EvidenceMaximisation

results, opt_params = EvidenceMaximisation(
    model=model,
    ns_kwargs=dict(max_samples=1e5)
).train(num_steps=10)

Step 0: log Z = -0.7434 +- 0.0362:   0%|          | 0/10 [00:01<?, ?it/s]
Training Progress:   0%|          | 0/100 [00:00<?, ?it/s]
Step 0: log Z = -0.7434, l_oo = {'~': {'bias_scale_param': Array(0., dtype=float32), 'flat_params_param': Array(8.173739e-11, dtype=float32)}}:   0%|          | 0/100 [00:00<?, ?it/s]
Step 0: log Z = -0.7434, l_oo = {'~': {'bias_scale_param': Array(0., dtype=float32), 'flat_params_param': Array(8.173739e-11, dtype=float32)}}:   0%|          | 0/100 [00:00<?, ?it/s]
Step 0: log Z = -0.7434 +- 0.0362:  10%|█         | 1/10 [00:02<00:18,  2.03s/it]
Terminating at step 0 due to l(inf) small enough.
Step 1: log Z = -0.7095 +- 0.0159:  10%|█         | 1/10 [00:04<00:18,  2.03s/it]
Training Progress:   0%|          | 0/100 [00:00<?, ?it/s]
Step 0: log Z = -0.7095, l_oo = {'~': {'bias_scale_param': Array(8.951856e-10, dtype=float32), 'flat_params_param': Array(0.000218, dtype=float32)}}:   0%|          | 0/100 [00:01<?, ?it/s]
Step 0: log Z = -0.7095, l_oo = {'~': {'bias_scale_param': Array(8.951856e-10, dtype=float32), 'flat_params_param': Array(0.000218, dtype=float32)}}:   0%|          | 0/100 [00:01<?, ?it/s]
Step 0: log Z = -0.7095, l_oo = {'~': {'bias_scale_param': Array(8.951856e-10, dtype=float32), 'flat_params_param': Array(0.000218, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 1: log Z = -0.7103, l_oo = {'~': {'bias_scale_param': Array(2.5248637e-05, dtype=float32), 'flat_params_param': Array(0.08825424, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 1: log Z = -0.7103, l_oo = {'~': {'bias_scale_param': Array(2.5248637e-05, dtype=float32), 'flat_params_param': Array(0.08825424, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 2: log Z = -0.7096, l_oo = {'~': {'bias_scale_param': Array(2.0298241e-06, dtype=float32), 'flat_params_param': Array(0.02275545, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 2: log Z = -0.7096, l_oo = {'~': {'bias_scale_param': Array(2.0298241e-06, dtype=float32), 'flat_params_param': Array(0.02275545, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 3: log Z = -0.7097, l_oo = {'~': {'bias_scale_param': Array(7.4498766e-06, dtype=float32), 'flat_params_param': Array(0.04135855, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 3: log Z = -0.7097, l_oo = {'~': {'bias_scale_param': Array(7.4498766e-06, dtype=float32), 'flat_params_param': Array(0.04135855, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 4: log Z = -0.7100, l_oo = {'~': {'bias_scale_param': Array(1.9306628e-05, dtype=float32), 'flat_params_param': Array(0.06765009, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 4: log Z = -0.7100, l_oo = {'~': {'bias_scale_param': Array(1.9306628e-05, dtype=float32), 'flat_params_param': Array(0.06765009, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 5: log Z = -0.7099, l_oo = {'~': {'bias_scale_param': Array(1.1797916e-05, dtype=float32), 'flat_params_param': Array(0.05827373, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 5: log Z = -0.7099, l_oo = {'~': {'bias_scale_param': Array(1.1797916e-05, dtype=float32), 'flat_params_param': Array(0.05827373, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 6: log Z = -0.7096, l_oo = {'~': {'bias_scale_param': Array(1.3505763e-06, dtype=float32), 'flat_params_param': Array(0.03061554, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 6: log Z = -0.7096, l_oo = {'~': {'bias_scale_param': Array(1.3505763e-06, dtype=float32), 'flat_params_param': Array(0.03061554, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 7: log Z = -0.7095, l_oo = {'~': {'bias_scale_param': Array(3.7753918e-07, dtype=float32), 'flat_params_param': Array(0.00392881, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 7: log Z = -0.7095, l_oo = {'~': {'bias_scale_param': Array(3.7753918e-07, dtype=float32), 'flat_params_param': Array(0.00392881, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 8: log Z = -0.7096, l_oo = {'~': {'bias_scale_param': Array(7.0559586e-06, dtype=float32), 'flat_params_param': Array(0.03065771, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 8: log Z = -0.7096, l_oo = {'~': {'bias_scale_param': Array(7.0559586e-06, dtype=float32), 'flat_params_param': Array(0.03065771, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 9: log Z = -0.7097, l_oo = {'~': {'bias_scale_param': Array(1.2818397e-05, dtype=float32), 'flat_params_param': Array(0.04495464, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 9: log Z = -0.7097, l_oo = {'~': {'bias_scale_param': Array(1.2818397e-05, dtype=float32), 'flat_params_param': Array(0.04495464, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 10: log Z = -0.7097, l_oo = {'~': {'bias_scale_param': Array(1.2776985e-05, dtype=float32), 'flat_params_param': Array(0.04373395, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 10: log Z = -0.7097, l_oo = {'~': {'bias_scale_param': Array(1.2776985e-05, dtype=float32), 'flat_params_param': Array(0.04373395, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 11: log Z = -0.7096, l_oo = {'~': {'bias_scale_param': Array(7.536659e-06, dtype=float32), 'flat_params_param': Array(0.03027816, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 11: log Z = -0.7096, l_oo = {'~': {'bias_scale_param': Array(7.536659e-06, dtype=float32), 'flat_params_param': Array(0.03027816, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 12: log Z = -0.7095, l_oo = {'~': {'bias_scale_param': Array(6.605037e-08, dtype=float32), 'flat_params_param': Array(0.01042769, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 12: log Z = -0.7095, l_oo = {'~': {'bias_scale_param': Array(6.605037e-08, dtype=float32), 'flat_params_param': Array(0.01042769, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 13: log Z = -0.7095, l_oo = {'~': {'bias_scale_param': Array(7.225054e-06, dtype=float32), 'flat_params_param': Array(0.01181107, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 13: log Z = -0.7095, l_oo = {'~': {'bias_scale_param': Array(7.225054e-06, dtype=float32), 'flat_params_param': Array(0.01181107, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 14: log Z = -0.7096, l_oo = {'~': {'bias_scale_param': Array(1.4632719e-05, dtype=float32), 'flat_params_param': Array(0.02865441, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 14: log Z = -0.7096, l_oo = {'~': {'bias_scale_param': Array(1.4632719e-05, dtype=float32), 'flat_params_param': Array(0.02865441, dtype=float32)}}:   1%|          | 1/100 [00:01<02:02,  1.24s/it]
Step 14: log Z = -0.7096, l_oo = {'~': {'bias_scale_param': Array(1.4632719e-05, dtype=float32), 'flat_params_param': Array(0.02865441, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 15: log Z = -0.7096, l_oo = {'~': {'bias_scale_param': Array(2.4275681e-05, dtype=float32), 'flat_params_param': Array(0.0365797, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 15: log Z = -0.7096, l_oo = {'~': {'bias_scale_param': Array(2.4275681e-05, dtype=float32), 'flat_params_param': Array(0.0365797, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 16: log Z = -0.7096, l_oo = {'~': {'bias_scale_param': Array(3.6491398e-05, dtype=float32), 'flat_params_param': Array(0.03463286, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 16: log Z = -0.7096, l_oo = {'~': {'bias_scale_param': Array(3.6491398e-05, dtype=float32), 'flat_params_param': Array(0.03463286, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 17: log Z = -0.7095, l_oo = {'~': {'bias_scale_param': Array(4.932694e-05, dtype=float32), 'flat_params_param': Array(0.02478584, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 17: log Z = -0.7095, l_oo = {'~': {'bias_scale_param': Array(4.932694e-05, dtype=float32), 'flat_params_param': Array(0.02478584, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 18: log Z = -0.7094, l_oo = {'~': {'bias_scale_param': Array(6.206106e-05, dtype=float32), 'flat_params_param': Array(0.01070633, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 18: log Z = -0.7094, l_oo = {'~': {'bias_scale_param': Array(6.206106e-05, dtype=float32), 'flat_params_param': Array(0.01070633, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 19: log Z = -0.7094, l_oo = {'~': {'bias_scale_param': Array(7.869372e-05, dtype=float32), 'flat_params_param': Array(0.01980873, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 19: log Z = -0.7094, l_oo = {'~': {'bias_scale_param': Array(7.869372e-05, dtype=float32), 'flat_params_param': Array(0.01980873, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 20: log Z = -0.7093, l_oo = {'~': {'bias_scale_param': Array(0.00010865, dtype=float32), 'flat_params_param': Array(0.03086925, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 20: log Z = -0.7093, l_oo = {'~': {'bias_scale_param': Array(0.00010865, dtype=float32), 'flat_params_param': Array(0.03086925, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 21: log Z = -0.7092, l_oo = {'~': {'bias_scale_param': Array(0.00016444, dtype=float32), 'flat_params_param': Array(0.03915163, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 21: log Z = -0.7092, l_oo = {'~': {'bias_scale_param': Array(0.00016444, dtype=float32), 'flat_params_param': Array(0.03915163, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 22: log Z = -0.7089, l_oo = {'~': {'bias_scale_param': Array(0.00025797, dtype=float32), 'flat_params_param': Array(0.04473871, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 22: log Z = -0.7089, l_oo = {'~': {'bias_scale_param': Array(0.00025797, dtype=float32), 'flat_params_param': Array(0.04473871, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 23: log Z = -0.7085, l_oo = {'~': {'bias_scale_param': Array(0.00039705, dtype=float32), 'flat_params_param': Array(0.04846028, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 23: log Z = -0.7085, l_oo = {'~': {'bias_scale_param': Array(0.00039705, dtype=float32), 'flat_params_param': Array(0.04846028, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 24: log Z = -0.7080, l_oo = {'~': {'bias_scale_param': Array(0.00058304, dtype=float32), 'flat_params_param': Array(0.05128504, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 24: log Z = -0.7080, l_oo = {'~': {'bias_scale_param': Array(0.00058304, dtype=float32), 'flat_params_param': Array(0.05128504, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 25: log Z = -0.7073, l_oo = {'~': {'bias_scale_param': Array(0.00081098, dtype=float32), 'flat_params_param': Array(0.05403852, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 25: log Z = -0.7073, l_oo = {'~': {'bias_scale_param': Array(0.00081098, dtype=float32), 'flat_params_param': Array(0.05403852, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 26: log Z = -0.7064, l_oo = {'~': {'bias_scale_param': Array(0.00107132, dtype=float32), 'flat_params_param': Array(0.0572683, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 26: log Z = -0.7064, l_oo = {'~': {'bias_scale_param': Array(0.00107132, dtype=float32), 'flat_params_param': Array(0.0572683, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 27: log Z = -0.7054, l_oo = {'~': {'bias_scale_param': Array(0.00135123, dtype=float32), 'flat_params_param': Array(0.06116017, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 27: log Z = -0.7054, l_oo = {'~': {'bias_scale_param': Array(0.00135123, dtype=float32), 'flat_params_param': Array(0.06116017, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 28: log Z = -0.7042, l_oo = {'~': {'bias_scale_param': Array(0.00163456, dtype=float32), 'flat_params_param': Array(0.06555928, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 28: log Z = -0.7042, l_oo = {'~': {'bias_scale_param': Array(0.00163456, dtype=float32), 'flat_params_param': Array(0.06555928, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 29: log Z = -0.7029, l_oo = {'~': {'bias_scale_param': Array(0.00189997, dtype=float32), 'flat_params_param': Array(0.07004993, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 29: log Z = -0.7029, l_oo = {'~': {'bias_scale_param': Array(0.00189997, dtype=float32), 'flat_params_param': Array(0.07004993, dtype=float32)}}:  15%|█▌        | 15/100 [00:01<00:05, 15.20it/s]
Step 29: log Z = -0.7029, l_oo = {'~': {'bias_scale_param': Array(0.00189997, dtype=float32), 'flat_params_param': Array(0.07004993, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 30: log Z = -0.7016, l_oo = {'~': {'bias_scale_param': Array(0.00211941, dtype=float32), 'flat_params_param': Array(0.07410987, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 30: log Z = -0.7016, l_oo = {'~': {'bias_scale_param': Array(0.00211941, dtype=float32), 'flat_params_param': Array(0.07410987, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 31: log Z = -0.7002, l_oo = {'~': {'bias_scale_param': Array(0.00225954, dtype=float32), 'flat_params_param': Array(0.07719464, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 31: log Z = -0.7002, l_oo = {'~': {'bias_scale_param': Array(0.00225954, dtype=float32), 'flat_params_param': Array(0.07719464, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 32: log Z = -0.6988, l_oo = {'~': {'bias_scale_param': Array(0.0022888, dtype=float32), 'flat_params_param': Array(0.07885929, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 32: log Z = -0.6988, l_oo = {'~': {'bias_scale_param': Array(0.0022888, dtype=float32), 'flat_params_param': Array(0.07885929, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 33: log Z = -0.6975, l_oo = {'~': {'bias_scale_param': Array(0.00218866, dtype=float32), 'flat_params_param': Array(0.07882264, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 33: log Z = -0.6975, l_oo = {'~': {'bias_scale_param': Array(0.00218866, dtype=float32), 'flat_params_param': Array(0.07882264, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 34: log Z = -0.6963, l_oo = {'~': {'bias_scale_param': Array(0.00196553, dtype=float32), 'flat_params_param': Array(0.07703226, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 34: log Z = -0.6963, l_oo = {'~': {'bias_scale_param': Array(0.00196553, dtype=float32), 'flat_params_param': Array(0.07703226, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 35: log Z = -0.6952, l_oo = {'~': {'bias_scale_param': Array(0.0016574, dtype=float32), 'flat_params_param': Array(0.07368967, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 35: log Z = -0.6952, l_oo = {'~': {'bias_scale_param': Array(0.0016574, dtype=float32), 'flat_params_param': Array(0.07368967, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 36: log Z = -0.6942, l_oo = {'~': {'bias_scale_param': Array(0.00132996, dtype=float32), 'flat_params_param': Array(0.06917904, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 36: log Z = -0.6942, l_oo = {'~': {'bias_scale_param': Array(0.00132996, dtype=float32), 'flat_params_param': Array(0.06917904, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 37: log Z = -0.6933, l_oo = {'~': {'bias_scale_param': Array(0.00105867, dtype=float32), 'flat_params_param': Array(0.06399248, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 37: log Z = -0.6933, l_oo = {'~': {'bias_scale_param': Array(0.00105867, dtype=float32), 'flat_params_param': Array(0.06399248, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 38: log Z = -0.6926, l_oo = {'~': {'bias_scale_param': Array(0.00090383, dtype=float32), 'flat_params_param': Array(0.05862297, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 38: log Z = -0.6926, l_oo = {'~': {'bias_scale_param': Array(0.00090383, dtype=float32), 'flat_params_param': Array(0.05862297, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 39: log Z = -0.6920, l_oo = {'~': {'bias_scale_param': Array(0.00089213, dtype=float32), 'flat_params_param': Array(0.05348162, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 39: log Z = -0.6920, l_oo = {'~': {'bias_scale_param': Array(0.00089213, dtype=float32), 'flat_params_param': Array(0.05348162, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 40: log Z = -0.6915, l_oo = {'~': {'bias_scale_param': Array(0.00101556, dtype=float32), 'flat_params_param': Array(0.04884745, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 40: log Z = -0.6915, l_oo = {'~': {'bias_scale_param': Array(0.00101556, dtype=float32), 'flat_params_param': Array(0.04884745, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 41: log Z = -0.6910, l_oo = {'~': {'bias_scale_param': Array(0.00124473, dtype=float32), 'flat_params_param': Array(0.04485072, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 41: log Z = -0.6910, l_oo = {'~': {'bias_scale_param': Array(0.00124473, dtype=float32), 'flat_params_param': Array(0.04485072, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 42: log Z = -0.6905, l_oo = {'~': {'bias_scale_param': Array(0.00154522, dtype=float32), 'flat_params_param': Array(0.0414849, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 42: log Z = -0.6905, l_oo = {'~': {'bias_scale_param': Array(0.00154522, dtype=float32), 'flat_params_param': Array(0.0414849, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 43: log Z = -0.6901, l_oo = {'~': {'bias_scale_param': Array(0.00188929, dtype=float32), 'flat_params_param': Array(0.03866204, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 43: log Z = -0.6901, l_oo = {'~': {'bias_scale_param': Array(0.00188929, dtype=float32), 'flat_params_param': Array(0.03866204, dtype=float32)}}:  30%|███       | 30/100 [00:01<00:02, 32.43it/s]
Step 43: log Z = -0.6901, l_oo = {'~': {'bias_scale_param': Array(0.00188929, dtype=float32), 'flat_params_param': Array(0.03866204, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 44: log Z = -0.6897, l_oo = {'~': {'bias_scale_param': Array(0.00226054, dtype=float32), 'flat_params_param': Array(0.03622958, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 44: log Z = -0.6897, l_oo = {'~': {'bias_scale_param': Array(0.00226054, dtype=float32), 'flat_params_param': Array(0.03622958, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 45: log Z = -0.6893, l_oo = {'~': {'bias_scale_param': Array(0.00265374, dtype=float32), 'flat_params_param': Array(0.03402929, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 45: log Z = -0.6893, l_oo = {'~': {'bias_scale_param': Array(0.00265374, dtype=float32), 'flat_params_param': Array(0.03402929, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 46: log Z = -0.6889, l_oo = {'~': {'bias_scale_param': Array(0.00307196, dtype=float32), 'flat_params_param': Array(0.03191207, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 46: log Z = -0.6889, l_oo = {'~': {'bias_scale_param': Array(0.00307196, dtype=float32), 'flat_params_param': Array(0.03191207, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 47: log Z = -0.6885, l_oo = {'~': {'bias_scale_param': Array(0.00352221, dtype=float32), 'flat_params_param': Array(0.02977038, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 47: log Z = -0.6885, l_oo = {'~': {'bias_scale_param': Array(0.00352221, dtype=float32), 'flat_params_param': Array(0.02977038, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 48: log Z = -0.6881, l_oo = {'~': {'bias_scale_param': Array(0.00401119, dtype=float32), 'flat_params_param': Array(0.02755281, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 48: log Z = -0.6881, l_oo = {'~': {'bias_scale_param': Array(0.00401119, dtype=float32), 'flat_params_param': Array(0.02755281, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 49: log Z = -0.6878, l_oo = {'~': {'bias_scale_param': Array(0.00454199, dtype=float32), 'flat_params_param': Array(0.0252637, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 49: log Z = -0.6878, l_oo = {'~': {'bias_scale_param': Array(0.00454199, dtype=float32), 'flat_params_param': Array(0.0252637, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 50: log Z = -0.6874, l_oo = {'~': {'bias_scale_param': Array(0.0051123, dtype=float32), 'flat_params_param': Array(0.0229558, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 50: log Z = -0.6874, l_oo = {'~': {'bias_scale_param': Array(0.0051123, dtype=float32), 'flat_params_param': Array(0.0229558, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 51: log Z = -0.6871, l_oo = {'~': {'bias_scale_param': Array(0.005714, dtype=float32), 'flat_params_param': Array(0.0207146, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 51: log Z = -0.6871, l_oo = {'~': {'bias_scale_param': Array(0.005714, dtype=float32), 'flat_params_param': Array(0.0207146, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 52: log Z = -0.6868, l_oo = {'~': {'bias_scale_param': Array(0.00633417, dtype=float32), 'flat_params_param': Array(0.01864337, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 52: log Z = -0.6868, l_oo = {'~': {'bias_scale_param': Array(0.00633417, dtype=float32), 'flat_params_param': Array(0.01864337, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 53: log Z = -0.6865, l_oo = {'~': {'bias_scale_param': Array(0.00695712, dtype=float32), 'flat_params_param': Array(0.01682884, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 53: log Z = -0.6865, l_oo = {'~': {'bias_scale_param': Array(0.00695712, dtype=float32), 'flat_params_param': Array(0.01682884, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 54: log Z = -0.6863, l_oo = {'~': {'bias_scale_param': Array(0.00756734, dtype=float32), 'flat_params_param': Array(0.01534054, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 54: log Z = -0.6863, l_oo = {'~': {'bias_scale_param': Array(0.00756734, dtype=float32), 'flat_params_param': Array(0.01534054, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 55: log Z = -0.6860, l_oo = {'~': {'bias_scale_param': Array(0.00815301, dtype=float32), 'flat_params_param': Array(0.01420399, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 55: log Z = -0.6860, l_oo = {'~': {'bias_scale_param': Array(0.00815301, dtype=float32), 'flat_params_param': Array(0.01420399, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 56: log Z = -0.6858, l_oo = {'~': {'bias_scale_param': Array(0.00870881, dtype=float32), 'flat_params_param': Array(0.01340893, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 56: log Z = -0.6858, l_oo = {'~': {'bias_scale_param': Array(0.00870881, dtype=float32), 'flat_params_param': Array(0.01340893, dtype=float32)}}:  44%|████▍     | 44/100 [00:01<00:01, 48.65it/s]
Step 56: log Z = -0.6858, l_oo = {'~': {'bias_scale_param': Array(0.00870881, dtype=float32), 'flat_params_param': Array(0.01340893, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 57: log Z = -0.6856, l_oo = {'~': {'bias_scale_param': Array(0.00923765, dtype=float32), 'flat_params_param': Array(0.01290459, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 57: log Z = -0.6856, l_oo = {'~': {'bias_scale_param': Array(0.00923765, dtype=float32), 'flat_params_param': Array(0.01290459, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 58: log Z = -0.6854, l_oo = {'~': {'bias_scale_param': Array(0.00975054, dtype=float32), 'flat_params_param': Array(0.01261501, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 58: log Z = -0.6854, l_oo = {'~': {'bias_scale_param': Array(0.00975054, dtype=float32), 'flat_params_param': Array(0.01261501, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 59: log Z = -0.6851, l_oo = {'~': {'bias_scale_param': Array(0.0102637, dtype=float32), 'flat_params_param': Array(0.01246421, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 59: log Z = -0.6851, l_oo = {'~': {'bias_scale_param': Array(0.0102637, dtype=float32), 'flat_params_param': Array(0.01246421, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 60: log Z = -0.6849, l_oo = {'~': {'bias_scale_param': Array(0.01079463, dtype=float32), 'flat_params_param': Array(0.01238908, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 60: log Z = -0.6849, l_oo = {'~': {'bias_scale_param': Array(0.01079463, dtype=float32), 'flat_params_param': Array(0.01238908, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 61: log Z = -0.6846, l_oo = {'~': {'bias_scale_param': Array(0.01135846, dtype=float32), 'flat_params_param': Array(0.01236261, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 61: log Z = -0.6846, l_oo = {'~': {'bias_scale_param': Array(0.01135846, dtype=float32), 'flat_params_param': Array(0.01236261, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 62: log Z = -0.6843, l_oo = {'~': {'bias_scale_param': Array(0.01196601, dtype=float32), 'flat_params_param': Array(0.01239257, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 62: log Z = -0.6843, l_oo = {'~': {'bias_scale_param': Array(0.01196601, dtype=float32), 'flat_params_param': Array(0.01239257, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 63: log Z = -0.6840, l_oo = {'~': {'bias_scale_param': Array(0.01262447, dtype=float32), 'flat_params_param': Array(0.01251375, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 63: log Z = -0.6840, l_oo = {'~': {'bias_scale_param': Array(0.01262447, dtype=float32), 'flat_params_param': Array(0.01251375, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 64: log Z = -0.6837, l_oo = {'~': {'bias_scale_param': Array(0.01333914, dtype=float32), 'flat_params_param': Array(0.0127714, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 64: log Z = -0.6837, l_oo = {'~': {'bias_scale_param': Array(0.01333914, dtype=float32), 'flat_params_param': Array(0.0127714, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 65: log Z = -0.6834, l_oo = {'~': {'bias_scale_param': Array(0.0141158, dtype=float32), 'flat_params_param': Array(0.01320754, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 65: log Z = -0.6834, l_oo = {'~': {'bias_scale_param': Array(0.0141158, dtype=float32), 'flat_params_param': Array(0.01320754, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 66: log Z = -0.6830, l_oo = {'~': {'bias_scale_param': Array(0.01496205, dtype=float32), 'flat_params_param': Array(0.0138431, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 66: log Z = -0.6830, l_oo = {'~': {'bias_scale_param': Array(0.01496205, dtype=float32), 'flat_params_param': Array(0.0138431, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 67: log Z = -0.6826, l_oo = {'~': {'bias_scale_param': Array(0.01588697, dtype=float32), 'flat_params_param': Array(0.01468201, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 67: log Z = -0.6826, l_oo = {'~': {'bias_scale_param': Array(0.01588697, dtype=float32), 'flat_params_param': Array(0.01468201, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 68: log Z = -0.6822, l_oo = {'~': {'bias_scale_param': Array(0.01689911, dtype=float32), 'flat_params_param': Array(0.0157062, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 68: log Z = -0.6822, l_oo = {'~': {'bias_scale_param': Array(0.01689911, dtype=float32), 'flat_params_param': Array(0.0157062, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 69: log Z = -0.6817, l_oo = {'~': {'bias_scale_param': Array(0.01800231, dtype=float32), 'flat_params_param': Array(0.01689135, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 69: log Z = -0.6817, l_oo = {'~': {'bias_scale_param': Array(0.01800231, dtype=float32), 'flat_params_param': Array(0.01689135, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 70: log Z = -0.6812, l_oo = {'~': {'bias_scale_param': Array(0.01918954, dtype=float32), 'flat_params_param': Array(0.01821935, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 70: log Z = -0.6812, l_oo = {'~': {'bias_scale_param': Array(0.01918954, dtype=float32), 'flat_params_param': Array(0.01821935, dtype=float32)}}:  57%|█████▋    | 57/100 [00:01<00:00, 62.51it/s]
Step 70: log Z = -0.6812, l_oo = {'~': {'bias_scale_param': Array(0.01918954, dtype=float32), 'flat_params_param': Array(0.01821935, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 71: log Z = -0.6806, l_oo = {'~': {'bias_scale_param': Array(0.02043372, dtype=float32), 'flat_params_param': Array(0.01968751, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 71: log Z = -0.6806, l_oo = {'~': {'bias_scale_param': Array(0.02043372, dtype=float32), 'flat_params_param': Array(0.01968751, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 72: log Z = -0.6800, l_oo = {'~': {'bias_scale_param': Array(0.02167499, dtype=float32), 'flat_params_param': Array(0.02132084, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 72: log Z = -0.6800, l_oo = {'~': {'bias_scale_param': Array(0.02167499, dtype=float32), 'flat_params_param': Array(0.02132084, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 73: log Z = -0.6792, l_oo = {'~': {'bias_scale_param': Array(0.02280507, dtype=float32), 'flat_params_param': Array(0.02317252, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 73: log Z = -0.6792, l_oo = {'~': {'bias_scale_param': Array(0.02280507, dtype=float32), 'flat_params_param': Array(0.02317252, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 74: log Z = -0.6784, l_oo = {'~': {'bias_scale_param': Array(0.02365118, dtype=float32), 'flat_params_param': Array(0.02531132, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 74: log Z = -0.6784, l_oo = {'~': {'bias_scale_param': Array(0.02365118, dtype=float32), 'flat_params_param': Array(0.02531132, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 75: log Z = -0.6775, l_oo = {'~': {'bias_scale_param': Array(0.0239659, dtype=float32), 'flat_params_param': Array(0.02779424, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 75: log Z = -0.6775, l_oo = {'~': {'bias_scale_param': Array(0.0239659, dtype=float32), 'flat_params_param': Array(0.02779424, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 76: log Z = -0.6765, l_oo = {'~': {'bias_scale_param': Array(0.02343715, dtype=float32), 'flat_params_param': Array(0.03062515, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 76: log Z = -0.6765, l_oo = {'~': {'bias_scale_param': Array(0.02343715, dtype=float32), 'flat_params_param': Array(0.03062515, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 77: log Z = -0.6755, l_oo = {'~': {'bias_scale_param': Array(0.0217507, dtype=float32), 'flat_params_param': Array(0.03370814, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 77: log Z = -0.6755, l_oo = {'~': {'bias_scale_param': Array(0.0217507, dtype=float32), 'flat_params_param': Array(0.03370814, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 78: log Z = -0.6746, l_oo = {'~': {'bias_scale_param': Array(0.01873951, dtype=float32), 'flat_params_param': Array(0.03681064, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 78: log Z = -0.6746, l_oo = {'~': {'bias_scale_param': Array(0.01873951, dtype=float32), 'flat_params_param': Array(0.03681064, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 79: log Z = -0.6739, l_oo = {'~': {'bias_scale_param': Array(0.01458043, dtype=float32), 'flat_params_param': Array(0.03958978, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 79: log Z = -0.6739, l_oo = {'~': {'bias_scale_param': Array(0.01458043, dtype=float32), 'flat_params_param': Array(0.03958978, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 80: log Z = -0.6734, l_oo = {'~': {'bias_scale_param': Array(0.00986312, dtype=float32), 'flat_params_param': Array(0.04171936, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 80: log Z = -0.6734, l_oo = {'~': {'bias_scale_param': Array(0.00986312, dtype=float32), 'flat_params_param': Array(0.04171936, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 81: log Z = -0.6731, l_oo = {'~': {'bias_scale_param': Array(0.00538962, dtype=float32), 'flat_params_param': Array(0.04303171, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 81: log Z = -0.6731, l_oo = {'~': {'bias_scale_param': Array(0.00538962, dtype=float32), 'flat_params_param': Array(0.04303171, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 82: log Z = -0.6730, l_oo = {'~': {'bias_scale_param': Array(0.00182056, dtype=float32), 'flat_params_param': Array(0.04359164, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 82: log Z = -0.6730, l_oo = {'~': {'bias_scale_param': Array(0.00182056, dtype=float32), 'flat_params_param': Array(0.04359164, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 83: log Z = -0.6728, l_oo = {'~': {'bias_scale_param': Array(0.00053212, dtype=float32), 'flat_params_param': Array(0.04360714, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 83: log Z = -0.6728, l_oo = {'~': {'bias_scale_param': Array(0.00053212, dtype=float32), 'flat_params_param': Array(0.04360714, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 84: log Z = -0.6725, l_oo = {'~': {'bias_scale_param': Array(0.0016527, dtype=float32), 'flat_params_param': Array(0.04329898, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 84: log Z = -0.6725, l_oo = {'~': {'bias_scale_param': Array(0.0016527, dtype=float32), 'flat_params_param': Array(0.04329898, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 85: log Z = -0.6721, l_oo = {'~': {'bias_scale_param': Array(0.00165525, dtype=float32), 'flat_params_param': Array(0.04281161, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 85: log Z = -0.6721, l_oo = {'~': {'bias_scale_param': Array(0.00165525, dtype=float32), 'flat_params_param': Array(0.04281161, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 77.69it/s]
Step 85: log Z = -0.6721, l_oo = {'~': {'bias_scale_param': Array(0.00165525, dtype=float32), 'flat_params_param': Array(0.04281161, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 86: log Z = -0.6715, l_oo = {'~': {'bias_scale_param': Array(0.00066903, dtype=float32), 'flat_params_param': Array(0.04219921, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 86: log Z = -0.6715, l_oo = {'~': {'bias_scale_param': Array(0.00066903, dtype=float32), 'flat_params_param': Array(0.04219921, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 87: log Z = -0.6708, l_oo = {'~': {'bias_scale_param': Array(0.00119286, dtype=float32), 'flat_params_param': Array(0.04145166, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 87: log Z = -0.6708, l_oo = {'~': {'bias_scale_param': Array(0.00119286, dtype=float32), 'flat_params_param': Array(0.04145166, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 88: log Z = -0.6700, l_oo = {'~': {'bias_scale_param': Array(0.00381271, dtype=float32), 'flat_params_param': Array(0.04052468, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 88: log Z = -0.6700, l_oo = {'~': {'bias_scale_param': Array(0.00381271, dtype=float32), 'flat_params_param': Array(0.04052468, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 89: log Z = -0.6692, l_oo = {'~': {'bias_scale_param': Array(0.00703213, dtype=float32), 'flat_params_param': Array(0.03937577, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 89: log Z = -0.6692, l_oo = {'~': {'bias_scale_param': Array(0.00703213, dtype=float32), 'flat_params_param': Array(0.03937577, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 90: log Z = -0.6685, l_oo = {'~': {'bias_scale_param': Array(0.01062686, dtype=float32), 'flat_params_param': Array(0.03800843, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 90: log Z = -0.6685, l_oo = {'~': {'bias_scale_param': Array(0.01062686, dtype=float32), 'flat_params_param': Array(0.03800843, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 91: log Z = -0.6679, l_oo = {'~': {'bias_scale_param': Array(0.01431472, dtype=float32), 'flat_params_param': Array(0.03648005, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 91: log Z = -0.6679, l_oo = {'~': {'bias_scale_param': Array(0.01431472, dtype=float32), 'flat_params_param': Array(0.03648005, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 92: log Z = -0.6674, l_oo = {'~': {'bias_scale_param': Array(0.01780222, dtype=float32), 'flat_params_param': Array(0.03489604, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 92: log Z = -0.6674, l_oo = {'~': {'bias_scale_param': Array(0.01780222, dtype=float32), 'flat_params_param': Array(0.03489604, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 93: log Z = -0.6669, l_oo = {'~': {'bias_scale_param': Array(0.02084826, dtype=float32), 'flat_params_param': Array(0.0334052, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 93: log Z = -0.6669, l_oo = {'~': {'bias_scale_param': Array(0.02084826, dtype=float32), 'flat_params_param': Array(0.0334052, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 94: log Z = -0.6664, l_oo = {'~': {'bias_scale_param': Array(0.02330641, dtype=float32), 'flat_params_param': Array(0.03214372, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 94: log Z = -0.6664, l_oo = {'~': {'bias_scale_param': Array(0.02330641, dtype=float32), 'flat_params_param': Array(0.03214372, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 95: log Z = -0.6659, l_oo = {'~': {'bias_scale_param': Array(0.02512472, dtype=float32), 'flat_params_param': Array(0.03120418, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 95: log Z = -0.6659, l_oo = {'~': {'bias_scale_param': Array(0.02512472, dtype=float32), 'flat_params_param': Array(0.03120418, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 96: log Z = -0.6654, l_oo = {'~': {'bias_scale_param': Array(0.02631235, dtype=float32), 'flat_params_param': Array(0.03062845, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 96: log Z = -0.6654, l_oo = {'~': {'bias_scale_param': Array(0.02631235, dtype=float32), 'flat_params_param': Array(0.03062845, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 97: log Z = -0.6648, l_oo = {'~': {'bias_scale_param': Array(0.02690648, dtype=float32), 'flat_params_param': Array(0.03041125, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 97: log Z = -0.6648, l_oo = {'~': {'bias_scale_param': Array(0.02690648, dtype=float32), 'flat_params_param': Array(0.03041125, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 98: log Z = -0.6642, l_oo = {'~': {'bias_scale_param': Array(0.0269525, dtype=float32), 'flat_params_param': Array(0.03050743, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 98: log Z = -0.6642, l_oo = {'~': {'bias_scale_param': Array(0.0269525, dtype=float32), 'flat_params_param': Array(0.03050743, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 99: log Z = -0.6635, l_oo = {'~': {'bias_scale_param': Array(0.02650591, dtype=float32), 'flat_params_param': Array(0.03084504, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 99: log Z = -0.6635, l_oo = {'~': {'bias_scale_param': Array(0.02650591, dtype=float32), 'flat_params_param': Array(0.03084504, dtype=float32)}}:  86%|████████▌ | 86/100 [00:01<00:00, 92.44it/s]
Step 99: log Z = -0.6635, l_oo = {'~': {'bias_scale_param': Array(0.02650591, dtype=float32), 'flat_params_param': Array(0.03084504, dtype=float32)}}: 100%|██████████| 100/100 [00:01<00:00, 50.86it/s]
Step 2: log Z = -0.6338 +- 0.0198:  20%|██        | 2/10 [00:07<00:26,  3.31s/it]
Training Progress:   0%|          | 0/100 [00:00<?, ?it/s]
Step 0: log Z = -0.6338, l_oo = {'~': {'bias_scale_param': Array(0.00283509, dtype=float32), 'flat_params_param': Array(0.09193294, dtype=float32)}}:   0%|          | 0/100 [00:00<?, ?it/s]
Step 0: log Z = -0.6338, l_oo = {'~': {'bias_scale_param': Array(0.00283509, dtype=float32), 'flat_params_param': Array(0.09193294, dtype=float32)}}:   0%|          | 0/100 [00:00<?, ?it/s]
Step 0: log Z = -0.6338, l_oo = {'~': {'bias_scale_param': Array(0.00283509, dtype=float32), 'flat_params_param': Array(0.09193294, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 1: log Z = -0.6315, l_oo = {'~': {'bias_scale_param': Array(0.00424997, dtype=float32), 'flat_params_param': Array(0.06334927, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 1: log Z = -0.6315, l_oo = {'~': {'bias_scale_param': Array(0.00424997, dtype=float32), 'flat_params_param': Array(0.06334927, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 2: log Z = -0.6298, l_oo = {'~': {'bias_scale_param': Array(0.01054184, dtype=float32), 'flat_params_param': Array(0.05907187, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 2: log Z = -0.6298, l_oo = {'~': {'bias_scale_param': Array(0.01054184, dtype=float32), 'flat_params_param': Array(0.05907187, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 3: log Z = -0.6286, l_oo = {'~': {'bias_scale_param': Array(0.01661834, dtype=float32), 'flat_params_param': Array(0.05784559, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 3: log Z = -0.6286, l_oo = {'~': {'bias_scale_param': Array(0.01661834, dtype=float32), 'flat_params_param': Array(0.05784559, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 4: log Z = -0.6279, l_oo = {'~': {'bias_scale_param': Array(0.02230016, dtype=float32), 'flat_params_param': Array(0.05628166, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 4: log Z = -0.6279, l_oo = {'~': {'bias_scale_param': Array(0.02230016, dtype=float32), 'flat_params_param': Array(0.05628166, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 5: log Z = -0.6274, l_oo = {'~': {'bias_scale_param': Array(0.0268275, dtype=float32), 'flat_params_param': Array(0.0547563, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 5: log Z = -0.6274, l_oo = {'~': {'bias_scale_param': Array(0.0268275, dtype=float32), 'flat_params_param': Array(0.0547563, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 6: log Z = -0.6269, l_oo = {'~': {'bias_scale_param': Array(0.02948412, dtype=float32), 'flat_params_param': Array(0.05375652, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 6: log Z = -0.6269, l_oo = {'~': {'bias_scale_param': Array(0.02948412, dtype=float32), 'flat_params_param': Array(0.05375652, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 7: log Z = -0.6263, l_oo = {'~': {'bias_scale_param': Array(0.03019085, dtype=float32), 'flat_params_param': Array(0.05340351, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 7: log Z = -0.6263, l_oo = {'~': {'bias_scale_param': Array(0.03019085, dtype=float32), 'flat_params_param': Array(0.05340351, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 8: log Z = -0.6255, l_oo = {'~': {'bias_scale_param': Array(0.02937021, dtype=float32), 'flat_params_param': Array(0.05346937, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 8: log Z = -0.6255, l_oo = {'~': {'bias_scale_param': Array(0.02937021, dtype=float32), 'flat_params_param': Array(0.05346937, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 9: log Z = -0.6246, l_oo = {'~': {'bias_scale_param': Array(0.02756785, dtype=float32), 'flat_params_param': Array(0.05366959, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 9: log Z = -0.6246, l_oo = {'~': {'bias_scale_param': Array(0.02756785, dtype=float32), 'flat_params_param': Array(0.05366959, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 10: log Z = -0.6236, l_oo = {'~': {'bias_scale_param': Array(0.02530158, dtype=float32), 'flat_params_param': Array(0.0537508, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 10: log Z = -0.6236, l_oo = {'~': {'bias_scale_param': Array(0.02530158, dtype=float32), 'flat_params_param': Array(0.0537508, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 11: log Z = -0.6228, l_oo = {'~': {'bias_scale_param': Array(0.02303453, dtype=float32), 'flat_params_param': Array(0.05355982, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 11: log Z = -0.6228, l_oo = {'~': {'bias_scale_param': Array(0.02303453, dtype=float32), 'flat_params_param': Array(0.05355982, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 12: log Z = -0.6220, l_oo = {'~': {'bias_scale_param': Array(0.02113809, dtype=float32), 'flat_params_param': Array(0.05304517, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 12: log Z = -0.6220, l_oo = {'~': {'bias_scale_param': Array(0.02113809, dtype=float32), 'flat_params_param': Array(0.05304517, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 13: log Z = -0.6213, l_oo = {'~': {'bias_scale_param': Array(0.01983949, dtype=float32), 'flat_params_param': Array(0.05224835, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 13: log Z = -0.6213, l_oo = {'~': {'bias_scale_param': Array(0.01983949, dtype=float32), 'flat_params_param': Array(0.05224835, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 14: log Z = -0.6206, l_oo = {'~': {'bias_scale_param': Array(0.01919314, dtype=float32), 'flat_params_param': Array(0.05129877, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 14: log Z = -0.6206, l_oo = {'~': {'bias_scale_param': Array(0.01919314, dtype=float32), 'flat_params_param': Array(0.05129877, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 15: log Z = -0.6200, l_oo = {'~': {'bias_scale_param': Array(0.01912256, dtype=float32), 'flat_params_param': Array(0.05035717, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 15: log Z = -0.6200, l_oo = {'~': {'bias_scale_param': Array(0.01912256, dtype=float32), 'flat_params_param': Array(0.05035717, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 16: log Z = -0.6193, l_oo = {'~': {'bias_scale_param': Array(0.01950298, dtype=float32), 'flat_params_param': Array(0.04955395, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 16: log Z = -0.6193, l_oo = {'~': {'bias_scale_param': Array(0.01950298, dtype=float32), 'flat_params_param': Array(0.04955395, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 17: log Z = -0.6186, l_oo = {'~': {'bias_scale_param': Array(0.02022335, dtype=float32), 'flat_params_param': Array(0.04893576, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 17: log Z = -0.6186, l_oo = {'~': {'bias_scale_param': Array(0.02022335, dtype=float32), 'flat_params_param': Array(0.04893576, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 18: log Z = -0.6178, l_oo = {'~': {'bias_scale_param': Array(0.02120227, dtype=float32), 'flat_params_param': Array(0.04850828, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 18: log Z = -0.6178, l_oo = {'~': {'bias_scale_param': Array(0.02120227, dtype=float32), 'flat_params_param': Array(0.04850828, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 19: log Z = -0.6170, l_oo = {'~': {'bias_scale_param': Array(0.02237946, dtype=float32), 'flat_params_param': Array(0.04822788, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 19: log Z = -0.6170, l_oo = {'~': {'bias_scale_param': Array(0.02237946, dtype=float32), 'flat_params_param': Array(0.04822788, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 20: log Z = -0.6162, l_oo = {'~': {'bias_scale_param': Array(0.02370177, dtype=float32), 'flat_params_param': Array(0.0480481, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 20: log Z = -0.6162, l_oo = {'~': {'bias_scale_param': Array(0.02370177, dtype=float32), 'flat_params_param': Array(0.0480481, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 21: log Z = -0.6154, l_oo = {'~': {'bias_scale_param': Array(0.02511367, dtype=float32), 'flat_params_param': Array(0.04790993, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 21: log Z = -0.6154, l_oo = {'~': {'bias_scale_param': Array(0.02511367, dtype=float32), 'flat_params_param': Array(0.04790993, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 22: log Z = -0.6147, l_oo = {'~': {'bias_scale_param': Array(0.02655116, dtype=float32), 'flat_params_param': Array(0.04776876, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 22: log Z = -0.6147, l_oo = {'~': {'bias_scale_param': Array(0.02655116, dtype=float32), 'flat_params_param': Array(0.04776876, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 23: log Z = -0.6140, l_oo = {'~': {'bias_scale_param': Array(0.02793453, dtype=float32), 'flat_params_param': Array(0.04758329, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 23: log Z = -0.6140, l_oo = {'~': {'bias_scale_param': Array(0.02793453, dtype=float32), 'flat_params_param': Array(0.04758329, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 24: log Z = -0.6133, l_oo = {'~': {'bias_scale_param': Array(0.02917, dtype=float32), 'flat_params_param': Array(0.04734072, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 24: log Z = -0.6133, l_oo = {'~': {'bias_scale_param': Array(0.02917, dtype=float32), 'flat_params_param': Array(0.04734072, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 25: log Z = -0.6126, l_oo = {'~': {'bias_scale_param': Array(0.03015595, dtype=float32), 'flat_params_param': Array(0.04704747, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 25: log Z = -0.6126, l_oo = {'~': {'bias_scale_param': Array(0.03015595, dtype=float32), 'flat_params_param': Array(0.04704747, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 26: log Z = -0.6118, l_oo = {'~': {'bias_scale_param': Array(0.03080537, dtype=float32), 'flat_params_param': Array(0.0467317, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 26: log Z = -0.6118, l_oo = {'~': {'bias_scale_param': Array(0.03080537, dtype=float32), 'flat_params_param': Array(0.0467317, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 27: log Z = -0.6111, l_oo = {'~': {'bias_scale_param': Array(0.03107689, dtype=float32), 'flat_params_param': Array(0.04642289, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 27: log Z = -0.6111, l_oo = {'~': {'bias_scale_param': Array(0.03107689, dtype=float32), 'flat_params_param': Array(0.04642289, dtype=float32)}}:   1%|          | 1/100 [00:00<01:01,  1.62it/s]
Step 27: log Z = -0.6111, l_oo = {'~': {'bias_scale_param': Array(0.03107689, dtype=float32), 'flat_params_param': Array(0.04642289, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 28: log Z = -0.6103, l_oo = {'~': {'bias_scale_param': Array(0.03099425, dtype=float32), 'flat_params_param': Array(0.04612824, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 28: log Z = -0.6103, l_oo = {'~': {'bias_scale_param': Array(0.03099425, dtype=float32), 'flat_params_param': Array(0.04612824, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 29: log Z = -0.6095, l_oo = {'~': {'bias_scale_param': Array(0.03064999, dtype=float32), 'flat_params_param': Array(0.04584248, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 29: log Z = -0.6095, l_oo = {'~': {'bias_scale_param': Array(0.03064999, dtype=float32), 'flat_params_param': Array(0.04584248, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 30: log Z = -0.6088, l_oo = {'~': {'bias_scale_param': Array(0.03018046, dtype=float32), 'flat_params_param': Array(0.04554508, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 30: log Z = -0.6088, l_oo = {'~': {'bias_scale_param': Array(0.03018046, dtype=float32), 'flat_params_param': Array(0.04554508, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 31: log Z = -0.6080, l_oo = {'~': {'bias_scale_param': Array(0.02973349, dtype=float32), 'flat_params_param': Array(0.04522373, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 31: log Z = -0.6080, l_oo = {'~': {'bias_scale_param': Array(0.02973349, dtype=float32), 'flat_params_param': Array(0.04522373, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 32: log Z = -0.6073, l_oo = {'~': {'bias_scale_param': Array(0.02943128, dtype=float32), 'flat_params_param': Array(0.04487236, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 32: log Z = -0.6073, l_oo = {'~': {'bias_scale_param': Array(0.02943128, dtype=float32), 'flat_params_param': Array(0.04487236, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 33: log Z = -0.6065, l_oo = {'~': {'bias_scale_param': Array(0.02934506, dtype=float32), 'flat_params_param': Array(0.04450695, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 33: log Z = -0.6065, l_oo = {'~': {'bias_scale_param': Array(0.02934506, dtype=float32), 'flat_params_param': Array(0.04450695, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 34: log Z = -0.6058, l_oo = {'~': {'bias_scale_param': Array(0.02949142, dtype=float32), 'flat_params_param': Array(0.04416151, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 34: log Z = -0.6058, l_oo = {'~': {'bias_scale_param': Array(0.02949142, dtype=float32), 'flat_params_param': Array(0.04416151, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 35: log Z = -0.6050, l_oo = {'~': {'bias_scale_param': Array(0.02984271, dtype=float32), 'flat_params_param': Array(0.04386447, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 35: log Z = -0.6050, l_oo = {'~': {'bias_scale_param': Array(0.02984271, dtype=float32), 'flat_params_param': Array(0.04386447, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 36: log Z = -0.6043, l_oo = {'~': {'bias_scale_param': Array(0.03035028, dtype=float32), 'flat_params_param': Array(0.04363741, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 36: log Z = -0.6043, l_oo = {'~': {'bias_scale_param': Array(0.03035028, dtype=float32), 'flat_params_param': Array(0.04363741, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 37: log Z = -0.6035, l_oo = {'~': {'bias_scale_param': Array(0.03096142, dtype=float32), 'flat_params_param': Array(0.04349267, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 37: log Z = -0.6035, l_oo = {'~': {'bias_scale_param': Array(0.03096142, dtype=float32), 'flat_params_param': Array(0.04349267, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 38: log Z = -0.6027, l_oo = {'~': {'bias_scale_param': Array(0.03163588, dtype=float32), 'flat_params_param': Array(0.04341535, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 38: log Z = -0.6027, l_oo = {'~': {'bias_scale_param': Array(0.03163588, dtype=float32), 'flat_params_param': Array(0.04341535, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 39: log Z = -0.6019, l_oo = {'~': {'bias_scale_param': Array(0.03234604, dtype=float32), 'flat_params_param': Array(0.04338428, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 39: log Z = -0.6019, l_oo = {'~': {'bias_scale_param': Array(0.03234604, dtype=float32), 'flat_params_param': Array(0.04338428, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 40: log Z = -0.6011, l_oo = {'~': {'bias_scale_param': Array(0.03307366, dtype=float32), 'flat_params_param': Array(0.04335999, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 40: log Z = -0.6011, l_oo = {'~': {'bias_scale_param': Array(0.03307366, dtype=float32), 'flat_params_param': Array(0.04335999, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 41: log Z = -0.6004, l_oo = {'~': {'bias_scale_param': Array(0.03379449, dtype=float32), 'flat_params_param': Array(0.04331908, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 41: log Z = -0.6004, l_oo = {'~': {'bias_scale_param': Array(0.03379449, dtype=float32), 'flat_params_param': Array(0.04331908, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 42: log Z = -0.5996, l_oo = {'~': {'bias_scale_param': Array(0.03446835, dtype=float32), 'flat_params_param': Array(0.04324327, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 42: log Z = -0.5996, l_oo = {'~': {'bias_scale_param': Array(0.03446835, dtype=float32), 'flat_params_param': Array(0.04324327, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 43: log Z = -0.5988, l_oo = {'~': {'bias_scale_param': Array(0.03504279, dtype=float32), 'flat_params_param': Array(0.04313416, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 43: log Z = -0.5988, l_oo = {'~': {'bias_scale_param': Array(0.03504279, dtype=float32), 'flat_params_param': Array(0.04313416, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 44: log Z = -0.5980, l_oo = {'~': {'bias_scale_param': Array(0.03546672, dtype=float32), 'flat_params_param': Array(0.04299977, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 44: log Z = -0.5980, l_oo = {'~': {'bias_scale_param': Array(0.03546672, dtype=float32), 'flat_params_param': Array(0.04299977, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 45: log Z = -0.5972, l_oo = {'~': {'bias_scale_param': Array(0.03571581, dtype=float32), 'flat_params_param': Array(0.04285891, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 45: log Z = -0.5972, l_oo = {'~': {'bias_scale_param': Array(0.03571581, dtype=float32), 'flat_params_param': Array(0.04285891, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 46: log Z = -0.5964, l_oo = {'~': {'bias_scale_param': Array(0.03580695, dtype=float32), 'flat_params_param': Array(0.04271338, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 46: log Z = -0.5964, l_oo = {'~': {'bias_scale_param': Array(0.03580695, dtype=float32), 'flat_params_param': Array(0.04271338, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 47: log Z = -0.5956, l_oo = {'~': {'bias_scale_param': Array(0.03579689, dtype=float32), 'flat_params_param': Array(0.04255684, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 47: log Z = -0.5956, l_oo = {'~': {'bias_scale_param': Array(0.03579689, dtype=float32), 'flat_params_param': Array(0.04255684, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 48: log Z = -0.5948, l_oo = {'~': {'bias_scale_param': Array(0.03576309, dtype=float32), 'flat_params_param': Array(0.04238421, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 48: log Z = -0.5948, l_oo = {'~': {'bias_scale_param': Array(0.03576309, dtype=float32), 'flat_params_param': Array(0.04238421, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 49: log Z = -0.5940, l_oo = {'~': {'bias_scale_param': Array(0.03577032, dtype=float32), 'flat_params_param': Array(0.04219769, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 49: log Z = -0.5940, l_oo = {'~': {'bias_scale_param': Array(0.03577032, dtype=float32), 'flat_params_param': Array(0.04219769, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 50: log Z = -0.5932, l_oo = {'~': {'bias_scale_param': Array(0.03585282, dtype=float32), 'flat_params_param': Array(0.0420054, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 50: log Z = -0.5932, l_oo = {'~': {'bias_scale_param': Array(0.03585282, dtype=float32), 'flat_params_param': Array(0.0420054, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 51: log Z = -0.5924, l_oo = {'~': {'bias_scale_param': Array(0.03600911, dtype=float32), 'flat_params_param': Array(0.04183294, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 51: log Z = -0.5924, l_oo = {'~': {'bias_scale_param': Array(0.03600911, dtype=float32), 'flat_params_param': Array(0.04183294, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 52: log Z = -0.5916, l_oo = {'~': {'bias_scale_param': Array(0.03622252, dtype=float32), 'flat_params_param': Array(0.04169576, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 52: log Z = -0.5916, l_oo = {'~': {'bias_scale_param': Array(0.03622252, dtype=float32), 'flat_params_param': Array(0.04169576, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 51.90it/s]
Step 52: log Z = -0.5916, l_oo = {'~': {'bias_scale_param': Array(0.03622252, dtype=float32), 'flat_params_param': Array(0.04169576, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 53: log Z = -0.5907, l_oo = {'~': {'bias_scale_param': Array(0.03647926, dtype=float32), 'flat_params_param': Array(0.04160981, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 53: log Z = -0.5907, l_oo = {'~': {'bias_scale_param': Array(0.03647926, dtype=float32), 'flat_params_param': Array(0.04160981, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 54: log Z = -0.5899, l_oo = {'~': {'bias_scale_param': Array(0.03678165, dtype=float32), 'flat_params_param': Array(0.04156455, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 54: log Z = -0.5899, l_oo = {'~': {'bias_scale_param': Array(0.03678165, dtype=float32), 'flat_params_param': Array(0.04156455, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 55: log Z = -0.5891, l_oo = {'~': {'bias_scale_param': Array(0.03714235, dtype=float32), 'flat_params_param': Array(0.04153764, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 55: log Z = -0.5891, l_oo = {'~': {'bias_scale_param': Array(0.03714235, dtype=float32), 'flat_params_param': Array(0.04153764, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 56: log Z = -0.5883, l_oo = {'~': {'bias_scale_param': Array(0.03756053, dtype=float32), 'flat_params_param': Array(0.04151739, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 56: log Z = -0.5883, l_oo = {'~': {'bias_scale_param': Array(0.03756053, dtype=float32), 'flat_params_param': Array(0.04151739, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 57: log Z = -0.5874, l_oo = {'~': {'bias_scale_param': Array(0.03801095, dtype=float32), 'flat_params_param': Array(0.04148965, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 57: log Z = -0.5874, l_oo = {'~': {'bias_scale_param': Array(0.03801095, dtype=float32), 'flat_params_param': Array(0.04148965, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 58: log Z = -0.5866, l_oo = {'~': {'bias_scale_param': Array(0.0384441, dtype=float32), 'flat_params_param': Array(0.04145376, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 58: log Z = -0.5866, l_oo = {'~': {'bias_scale_param': Array(0.0384441, dtype=float32), 'flat_params_param': Array(0.04145376, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 59: log Z = -0.5858, l_oo = {'~': {'bias_scale_param': Array(0.0388068, dtype=float32), 'flat_params_param': Array(0.04142112, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 59: log Z = -0.5858, l_oo = {'~': {'bias_scale_param': Array(0.0388068, dtype=float32), 'flat_params_param': Array(0.04142112, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 60: log Z = -0.5849, l_oo = {'~': {'bias_scale_param': Array(0.03907568, dtype=float32), 'flat_params_param': Array(0.04138831, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 60: log Z = -0.5849, l_oo = {'~': {'bias_scale_param': Array(0.03907568, dtype=float32), 'flat_params_param': Array(0.04138831, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 61: log Z = -0.5841, l_oo = {'~': {'bias_scale_param': Array(0.03926225, dtype=float32), 'flat_params_param': Array(0.04134691, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 61: log Z = -0.5841, l_oo = {'~': {'bias_scale_param': Array(0.03926225, dtype=float32), 'flat_params_param': Array(0.04134691, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 62: log Z = -0.5832, l_oo = {'~': {'bias_scale_param': Array(0.03940158, dtype=float32), 'flat_params_param': Array(0.04128166, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 62: log Z = -0.5832, l_oo = {'~': {'bias_scale_param': Array(0.03940158, dtype=float32), 'flat_params_param': Array(0.04128166, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 63: log Z = -0.5824, l_oo = {'~': {'bias_scale_param': Array(0.03952097, dtype=float32), 'flat_params_param': Array(0.04118687, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 63: log Z = -0.5824, l_oo = {'~': {'bias_scale_param': Array(0.03952097, dtype=float32), 'flat_params_param': Array(0.04118687, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 64: log Z = -0.5815, l_oo = {'~': {'bias_scale_param': Array(0.03962373, dtype=float32), 'flat_params_param': Array(0.04107518, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 64: log Z = -0.5815, l_oo = {'~': {'bias_scale_param': Array(0.03962373, dtype=float32), 'flat_params_param': Array(0.04107518, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 65: log Z = -0.5807, l_oo = {'~': {'bias_scale_param': Array(0.0396986, dtype=float32), 'flat_params_param': Array(0.04096446, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 65: log Z = -0.5807, l_oo = {'~': {'bias_scale_param': Array(0.0396986, dtype=float32), 'flat_params_param': Array(0.04096446, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 66: log Z = -0.5798, l_oo = {'~': {'bias_scale_param': Array(0.03974652, dtype=float32), 'flat_params_param': Array(0.04086997, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 66: log Z = -0.5798, l_oo = {'~': {'bias_scale_param': Array(0.03974652, dtype=float32), 'flat_params_param': Array(0.04086997, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 67: log Z = -0.5790, l_oo = {'~': {'bias_scale_param': Array(0.03979203, dtype=float32), 'flat_params_param': Array(0.04079261, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 67: log Z = -0.5790, l_oo = {'~': {'bias_scale_param': Array(0.03979203, dtype=float32), 'flat_params_param': Array(0.04079261, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 68: log Z = -0.5781, l_oo = {'~': {'bias_scale_param': Array(0.039875, dtype=float32), 'flat_params_param': Array(0.04072415, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 68: log Z = -0.5781, l_oo = {'~': {'bias_scale_param': Array(0.039875, dtype=float32), 'flat_params_param': Array(0.04072415, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 69: log Z = -0.5773, l_oo = {'~': {'bias_scale_param': Array(0.04001487, dtype=float32), 'flat_params_param': Array(0.04065831, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 69: log Z = -0.5773, l_oo = {'~': {'bias_scale_param': Array(0.04001487, dtype=float32), 'flat_params_param': Array(0.04065831, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 70: log Z = -0.5764, l_oo = {'~': {'bias_scale_param': Array(0.0401991, dtype=float32), 'flat_params_param': Array(0.04059947, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 70: log Z = -0.5764, l_oo = {'~': {'bias_scale_param': Array(0.0401991, dtype=float32), 'flat_params_param': Array(0.04059947, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 71: log Z = -0.5755, l_oo = {'~': {'bias_scale_param': Array(0.04039397, dtype=float32), 'flat_params_param': Array(0.04055861, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 71: log Z = -0.5755, l_oo = {'~': {'bias_scale_param': Array(0.04039397, dtype=float32), 'flat_params_param': Array(0.04055861, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 72: log Z = -0.5747, l_oo = {'~': {'bias_scale_param': Array(0.04057446, dtype=float32), 'flat_params_param': Array(0.04053985, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 72: log Z = -0.5747, l_oo = {'~': {'bias_scale_param': Array(0.04057446, dtype=float32), 'flat_params_param': Array(0.04053985, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 73: log Z = -0.5738, l_oo = {'~': {'bias_scale_param': Array(0.04074388, dtype=float32), 'flat_params_param': Array(0.0405279, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 73: log Z = -0.5738, l_oo = {'~': {'bias_scale_param': Array(0.04074388, dtype=float32), 'flat_params_param': Array(0.0405279, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 74: log Z = -0.5730, l_oo = {'~': {'bias_scale_param': Array(0.04091693, dtype=float32), 'flat_params_param': Array(0.04050782, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 74: log Z = -0.5730, l_oo = {'~': {'bias_scale_param': Array(0.04091693, dtype=float32), 'flat_params_param': Array(0.04050782, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 75: log Z = -0.5721, l_oo = {'~': {'bias_scale_param': Array(0.04109017, dtype=float32), 'flat_params_param': Array(0.0404711, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 75: log Z = -0.5721, l_oo = {'~': {'bias_scale_param': Array(0.04109017, dtype=float32), 'flat_params_param': Array(0.0404711, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 76: log Z = -0.5712, l_oo = {'~': {'bias_scale_param': Array(0.04123395, dtype=float32), 'flat_params_param': Array(0.04042754, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 76: log Z = -0.5712, l_oo = {'~': {'bias_scale_param': Array(0.04123395, dtype=float32), 'flat_params_param': Array(0.04042754, dtype=float32)}}:  53%|█████▎    | 53/100 [00:00<00:00, 93.27it/s]
Step 76: log Z = -0.5712, l_oo = {'~': {'bias_scale_param': Array(0.04123395, dtype=float32), 'flat_params_param': Array(0.04042754, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 77: log Z = -0.5704, l_oo = {'~': {'bias_scale_param': Array(0.04132401, dtype=float32), 'flat_params_param': Array(0.04038333, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 77: log Z = -0.5704, l_oo = {'~': {'bias_scale_param': Array(0.04132401, dtype=float32), 'flat_params_param': Array(0.04038333, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 78: log Z = -0.5695, l_oo = {'~': {'bias_scale_param': Array(0.04136629, dtype=float32), 'flat_params_param': Array(0.04033529, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 78: log Z = -0.5695, l_oo = {'~': {'bias_scale_param': Array(0.04136629, dtype=float32), 'flat_params_param': Array(0.04033529, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 79: log Z = -0.5686, l_oo = {'~': {'bias_scale_param': Array(0.04139062, dtype=float32), 'flat_params_param': Array(0.04027348, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 79: log Z = -0.5686, l_oo = {'~': {'bias_scale_param': Array(0.04139062, dtype=float32), 'flat_params_param': Array(0.04027348, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 80: log Z = -0.5678, l_oo = {'~': {'bias_scale_param': Array(0.04141472, dtype=float32), 'flat_params_param': Array(0.04019662, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 80: log Z = -0.5678, l_oo = {'~': {'bias_scale_param': Array(0.04141472, dtype=float32), 'flat_params_param': Array(0.04019662, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 81: log Z = -0.5669, l_oo = {'~': {'bias_scale_param': Array(0.04143006, dtype=float32), 'flat_params_param': Array(0.04011782, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 81: log Z = -0.5669, l_oo = {'~': {'bias_scale_param': Array(0.04143006, dtype=float32), 'flat_params_param': Array(0.04011782, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 82: log Z = -0.5660, l_oo = {'~': {'bias_scale_param': Array(0.04142382, dtype=float32), 'flat_params_param': Array(0.04004436, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 82: log Z = -0.5660, l_oo = {'~': {'bias_scale_param': Array(0.04142382, dtype=float32), 'flat_params_param': Array(0.04004436, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 83: log Z = -0.5652, l_oo = {'~': {'bias_scale_param': Array(0.04140533, dtype=float32), 'flat_params_param': Array(0.03998081, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 83: log Z = -0.5652, l_oo = {'~': {'bias_scale_param': Array(0.04140533, dtype=float32), 'flat_params_param': Array(0.03998081, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 84: log Z = -0.5643, l_oo = {'~': {'bias_scale_param': Array(0.04140439, dtype=float32), 'flat_params_param': Array(0.03991812, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 84: log Z = -0.5643, l_oo = {'~': {'bias_scale_param': Array(0.04140439, dtype=float32), 'flat_params_param': Array(0.03991812, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 85: log Z = -0.5634, l_oo = {'~': {'bias_scale_param': Array(0.0414287, dtype=float32), 'flat_params_param': Array(0.03985456, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 85: log Z = -0.5634, l_oo = {'~': {'bias_scale_param': Array(0.0414287, dtype=float32), 'flat_params_param': Array(0.03985456, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 86: log Z = -0.5626, l_oo = {'~': {'bias_scale_param': Array(0.04146137, dtype=float32), 'flat_params_param': Array(0.03979554, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 86: log Z = -0.5626, l_oo = {'~': {'bias_scale_param': Array(0.04146137, dtype=float32), 'flat_params_param': Array(0.03979554, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 87: log Z = -0.5617, l_oo = {'~': {'bias_scale_param': Array(0.04148152, dtype=float32), 'flat_params_param': Array(0.03974943, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 87: log Z = -0.5617, l_oo = {'~': {'bias_scale_param': Array(0.04148152, dtype=float32), 'flat_params_param': Array(0.03974943, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 88: log Z = -0.5608, l_oo = {'~': {'bias_scale_param': Array(0.04149666, dtype=float32), 'flat_params_param': Array(0.03970915, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 88: log Z = -0.5608, l_oo = {'~': {'bias_scale_param': Array(0.04149666, dtype=float32), 'flat_params_param': Array(0.03970915, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 89: log Z = -0.5600, l_oo = {'~': {'bias_scale_param': Array(0.04152056, dtype=float32), 'flat_params_param': Array(0.03966174, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 89: log Z = -0.5600, l_oo = {'~': {'bias_scale_param': Array(0.04152056, dtype=float32), 'flat_params_param': Array(0.03966174, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 90: log Z = -0.5591, l_oo = {'~': {'bias_scale_param': Array(0.04154646, dtype=float32), 'flat_params_param': Array(0.03960945, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 90: log Z = -0.5591, l_oo = {'~': {'bias_scale_param': Array(0.04154646, dtype=float32), 'flat_params_param': Array(0.03960945, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 91: log Z = -0.5583, l_oo = {'~': {'bias_scale_param': Array(0.04154824, dtype=float32), 'flat_params_param': Array(0.03955719, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 91: log Z = -0.5583, l_oo = {'~': {'bias_scale_param': Array(0.04154824, dtype=float32), 'flat_params_param': Array(0.03955719, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 92: log Z = -0.5574, l_oo = {'~': {'bias_scale_param': Array(0.04152125, dtype=float32), 'flat_params_param': Array(0.03950494, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 92: log Z = -0.5574, l_oo = {'~': {'bias_scale_param': Array(0.04152125, dtype=float32), 'flat_params_param': Array(0.03950494, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 93: log Z = -0.5566, l_oo = {'~': {'bias_scale_param': Array(0.04148177, dtype=float32), 'flat_params_param': Array(0.03944392, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 93: log Z = -0.5566, l_oo = {'~': {'bias_scale_param': Array(0.04148177, dtype=float32), 'flat_params_param': Array(0.03944392, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 94: log Z = -0.5557, l_oo = {'~': {'bias_scale_param': Array(0.04144024, dtype=float32), 'flat_params_param': Array(0.03937085, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 94: log Z = -0.5557, l_oo = {'~': {'bias_scale_param': Array(0.04144024, dtype=float32), 'flat_params_param': Array(0.03937085, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 95: log Z = -0.5549, l_oo = {'~': {'bias_scale_param': Array(0.04138011, dtype=float32), 'flat_params_param': Array(0.03929441, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 95: log Z = -0.5549, l_oo = {'~': {'bias_scale_param': Array(0.04138011, dtype=float32), 'flat_params_param': Array(0.03929441, dtype=float32)}}:  77%|███████▋  | 77/100 [00:00<00:00, 126.75it/s]
Step 96: log Z = -0.5540, l_oo = {'~': {'bias_scale_param': Array(0.04129523, dtype=float32), 'flat_params_param': Array(0.03921929, dtype=float32)}}:  77%|███████▋  | 77/100 [00:01<00:00, 126.75it/s]
Step 96: log Z = -0.5540, l_oo = {'~': {'bias_scale_param': Array(0.04129523, dtype=float32), 'flat_params_param': Array(0.03921929, dtype=float32)}}:  77%|███████▋  | 77/100 [00:01<00:00, 126.75it/s]
Step 97: log Z = -0.5532, l_oo = {'~': {'bias_scale_param': Array(0.04120345, dtype=float32), 'flat_params_param': Array(0.03914135, dtype=float32)}}:  77%|███████▋  | 77/100 [00:01<00:00, 126.75it/s]
Step 97: log Z = -0.5532, l_oo = {'~': {'bias_scale_param': Array(0.04120345, dtype=float32), 'flat_params_param': Array(0.03914135, dtype=float32)}}:  77%|███████▋  | 77/100 [00:01<00:00, 126.75it/s]
Step 98: log Z = -0.5523, l_oo = {'~': {'bias_scale_param': Array(0.04112121, dtype=float32), 'flat_params_param': Array(0.03905841, dtype=float32)}}:  77%|███████▋  | 77/100 [00:01<00:00, 126.75it/s]
Step 98: log Z = -0.5523, l_oo = {'~': {'bias_scale_param': Array(0.04112121, dtype=float32), 'flat_params_param': Array(0.03905841, dtype=float32)}}:  77%|███████▋  | 77/100 [00:01<00:00, 126.75it/s]
Step 99: log Z = -0.5515, l_oo = {'~': {'bias_scale_param': Array(0.04103538, dtype=float32), 'flat_params_param': Array(0.03897503, dtype=float32)}}:  77%|███████▋  | 77/100 [00:01<00:00, 126.75it/s]
Step 99: log Z = -0.5515, l_oo = {'~': {'bias_scale_param': Array(0.04103538, dtype=float32), 'flat_params_param': Array(0.03897503, dtype=float32)}}: 100%|██████████| 100/100 [00:01<00:00, 98.41it/s]
Step 3: log Z = -0.5743 +- 0.0295:  30%|███       | 3/10 [00:09<00:20,  2.92s/it]
Training Progress:   0%|          | 0/100 [00:00<?, ?it/s]
Step 0: log Z = -0.5743, l_oo = {'~': {'bias_scale_param': Array(0.04566116, dtype=float32), 'flat_params_param': Array(0.03655378, dtype=float32)}}:   0%|          | 0/100 [00:00<?, ?it/s]
Step 0: log Z = -0.5743, l_oo = {'~': {'bias_scale_param': Array(0.04566116, dtype=float32), 'flat_params_param': Array(0.03655378, dtype=float32)}}:   0%|          | 0/100 [00:00<?, ?it/s]
Step 0: log Z = -0.5743, l_oo = {'~': {'bias_scale_param': Array(0.04566116, dtype=float32), 'flat_params_param': Array(0.03655378, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 1: log Z = -0.5758, l_oo = {'~': {'bias_scale_param': Array(0.06274755, dtype=float32), 'flat_params_param': Array(0.10159741, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 1: log Z = -0.5758, l_oo = {'~': {'bias_scale_param': Array(0.06274755, dtype=float32), 'flat_params_param': Array(0.10159741, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 2: log Z = -0.5728, l_oo = {'~': {'bias_scale_param': Array(0.05001594, dtype=float32), 'flat_params_param': Array(0.03272018, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 2: log Z = -0.5728, l_oo = {'~': {'bias_scale_param': Array(0.05001594, dtype=float32), 'flat_params_param': Array(0.03272018, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 3: log Z = -0.5723, l_oo = {'~': {'bias_scale_param': Array(0.03871832, dtype=float32), 'flat_params_param': Array(0.06425725, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 3: log Z = -0.5723, l_oo = {'~': {'bias_scale_param': Array(0.03871832, dtype=float32), 'flat_params_param': Array(0.06425725, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 4: log Z = -0.5723, l_oo = {'~': {'bias_scale_param': Array(0.03383883, dtype=float32), 'flat_params_param': Array(0.09207651, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 4: log Z = -0.5723, l_oo = {'~': {'bias_scale_param': Array(0.03383883, dtype=float32), 'flat_params_param': Array(0.09207651, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 5: log Z = -0.5712, l_oo = {'~': {'bias_scale_param': Array(0.03436042, dtype=float32), 'flat_params_param': Array(0.07477449, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 5: log Z = -0.5712, l_oo = {'~': {'bias_scale_param': Array(0.03436042, dtype=float32), 'flat_params_param': Array(0.07477449, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 6: log Z = -0.5698, l_oo = {'~': {'bias_scale_param': Array(0.03790645, dtype=float32), 'flat_params_param': Array(0.03666956, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 6: log Z = -0.5698, l_oo = {'~': {'bias_scale_param': Array(0.03790645, dtype=float32), 'flat_params_param': Array(0.03666956, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 7: log Z = -0.5688, l_oo = {'~': {'bias_scale_param': Array(0.04272619, dtype=float32), 'flat_params_param': Array(0.0332367, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 7: log Z = -0.5688, l_oo = {'~': {'bias_scale_param': Array(0.04272619, dtype=float32), 'flat_params_param': Array(0.0332367, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 8: log Z = -0.5683, l_oo = {'~': {'bias_scale_param': Array(0.04710747, dtype=float32), 'flat_params_param': Array(0.06383491, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 8: log Z = -0.5683, l_oo = {'~': {'bias_scale_param': Array(0.04710747, dtype=float32), 'flat_params_param': Array(0.06383491, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 9: log Z = -0.5679, l_oo = {'~': {'bias_scale_param': Array(0.04945792, dtype=float32), 'flat_params_param': Array(0.08758664, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 9: log Z = -0.5679, l_oo = {'~': {'bias_scale_param': Array(0.04945792, dtype=float32), 'flat_params_param': Array(0.08758664, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 10: log Z = -0.5671, l_oo = {'~': {'bias_scale_param': Array(0.04920727, dtype=float32), 'flat_params_param': Array(0.08643748, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 10: log Z = -0.5671, l_oo = {'~': {'bias_scale_param': Array(0.04920727, dtype=float32), 'flat_params_param': Array(0.08643748, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 11: log Z = -0.5661, l_oo = {'~': {'bias_scale_param': Array(0.04691385, dtype=float32), 'flat_params_param': Array(0.06590208, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 11: log Z = -0.5661, l_oo = {'~': {'bias_scale_param': Array(0.04691385, dtype=float32), 'flat_params_param': Array(0.06590208, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 12: log Z = -0.5651, l_oo = {'~': {'bias_scale_param': Array(0.04363359, dtype=float32), 'flat_params_param': Array(0.03498027, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 12: log Z = -0.5651, l_oo = {'~': {'bias_scale_param': Array(0.04363359, dtype=float32), 'flat_params_param': Array(0.03498027, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 13: log Z = -0.5643, l_oo = {'~': {'bias_scale_param': Array(0.04050127, dtype=float32), 'flat_params_param': Array(0.03355434, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 13: log Z = -0.5643, l_oo = {'~': {'bias_scale_param': Array(0.04050127, dtype=float32), 'flat_params_param': Array(0.03355434, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 14: log Z = -0.5637, l_oo = {'~': {'bias_scale_param': Array(0.03846414, dtype=float32), 'flat_params_param': Array(0.03483785, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 14: log Z = -0.5637, l_oo = {'~': {'bias_scale_param': Array(0.03846414, dtype=float32), 'flat_params_param': Array(0.03483785, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 15: log Z = -0.5630, l_oo = {'~': {'bias_scale_param': Array(0.03803432, dtype=float32), 'flat_params_param': Array(0.03525149, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 15: log Z = -0.5630, l_oo = {'~': {'bias_scale_param': Array(0.03803432, dtype=float32), 'flat_params_param': Array(0.03525149, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 16: log Z = -0.5621, l_oo = {'~': {'bias_scale_param': Array(0.03922561, dtype=float32), 'flat_params_param': Array(0.03477792, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 16: log Z = -0.5621, l_oo = {'~': {'bias_scale_param': Array(0.03922561, dtype=float32), 'flat_params_param': Array(0.03477792, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 17: log Z = -0.5611, l_oo = {'~': {'bias_scale_param': Array(0.04166977, dtype=float32), 'flat_params_param': Array(0.03358628, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 17: log Z = -0.5611, l_oo = {'~': {'bias_scale_param': Array(0.04166977, dtype=float32), 'flat_params_param': Array(0.03358628, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 18: log Z = -0.5603, l_oo = {'~': {'bias_scale_param': Array(0.0447236, dtype=float32), 'flat_params_param': Array(0.03199239, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 18: log Z = -0.5603, l_oo = {'~': {'bias_scale_param': Array(0.0447236, dtype=float32), 'flat_params_param': Array(0.03199239, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 19: log Z = -0.5596, l_oo = {'~': {'bias_scale_param': Array(0.04753585, dtype=float32), 'flat_params_param': Array(0.03047284, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 19: log Z = -0.5596, l_oo = {'~': {'bias_scale_param': Array(0.04753585, dtype=float32), 'flat_params_param': Array(0.03047284, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 20: log Z = -0.5589, l_oo = {'~': {'bias_scale_param': Array(0.04926141, dtype=float32), 'flat_params_param': Array(0.03800664, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 20: log Z = -0.5589, l_oo = {'~': {'bias_scale_param': Array(0.04926141, dtype=float32), 'flat_params_param': Array(0.03800664, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 21: log Z = -0.5582, l_oo = {'~': {'bias_scale_param': Array(0.04943873, dtype=float32), 'flat_params_param': Array(0.03542551, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 21: log Z = -0.5582, l_oo = {'~': {'bias_scale_param': Array(0.04943873, dtype=float32), 'flat_params_param': Array(0.03542551, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 22: log Z = -0.5573, l_oo = {'~': {'bias_scale_param': Array(0.0482088, dtype=float32), 'flat_params_param': Array(0.03030985, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 22: log Z = -0.5573, l_oo = {'~': {'bias_scale_param': Array(0.0482088, dtype=float32), 'flat_params_param': Array(0.03030985, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 23: log Z = -0.5565, l_oo = {'~': {'bias_scale_param': Array(0.04618644, dtype=float32), 'flat_params_param': Array(0.03151589, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 23: log Z = -0.5565, l_oo = {'~': {'bias_scale_param': Array(0.04618644, dtype=float32), 'flat_params_param': Array(0.03151589, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 24: log Z = -0.5557, l_oo = {'~': {'bias_scale_param': Array(0.04417642, dtype=float32), 'flat_params_param': Array(0.03263777, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 24: log Z = -0.5557, l_oo = {'~': {'bias_scale_param': Array(0.04417642, dtype=float32), 'flat_params_param': Array(0.03263777, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 25: log Z = -0.5550, l_oo = {'~': {'bias_scale_param': Array(0.04288676, dtype=float32), 'flat_params_param': Array(0.03327743, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 25: log Z = -0.5550, l_oo = {'~': {'bias_scale_param': Array(0.04288676, dtype=float32), 'flat_params_param': Array(0.03327743, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 26: log Z = -0.5543, l_oo = {'~': {'bias_scale_param': Array(0.04269284, dtype=float32), 'flat_params_param': Array(0.03323334, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 26: log Z = -0.5543, l_oo = {'~': {'bias_scale_param': Array(0.04269284, dtype=float32), 'flat_params_param': Array(0.03323334, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 27: log Z = -0.5535, l_oo = {'~': {'bias_scale_param': Array(0.04354993, dtype=float32), 'flat_params_param': Array(0.03254041, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 27: log Z = -0.5535, l_oo = {'~': {'bias_scale_param': Array(0.04354993, dtype=float32), 'flat_params_param': Array(0.03254041, dtype=float32)}}:   1%|          | 1/100 [00:00<00:57,  1.72it/s]
Step 27: log Z = -0.5535, l_oo = {'~': {'bias_scale_param': Array(0.04354993, dtype=float32), 'flat_params_param': Array(0.03254041, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 28: log Z = -0.5527, l_oo = {'~': {'bias_scale_param': Array(0.04505412, dtype=float32), 'flat_params_param': Array(0.03142546, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 28: log Z = -0.5527, l_oo = {'~': {'bias_scale_param': Array(0.04505412, dtype=float32), 'flat_params_param': Array(0.03142546, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 29: log Z = -0.5519, l_oo = {'~': {'bias_scale_param': Array(0.04657435, dtype=float32), 'flat_params_param': Array(0.03025321, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 29: log Z = -0.5519, l_oo = {'~': {'bias_scale_param': Array(0.04657435, dtype=float32), 'flat_params_param': Array(0.03025321, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 30: log Z = -0.5512, l_oo = {'~': {'bias_scale_param': Array(0.047465, dtype=float32), 'flat_params_param': Array(0.02942213, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 30: log Z = -0.5512, l_oo = {'~': {'bias_scale_param': Array(0.047465, dtype=float32), 'flat_params_param': Array(0.02942213, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 31: log Z = -0.5505, l_oo = {'~': {'bias_scale_param': Array(0.04734327, dtype=float32), 'flat_params_param': Array(0.02919021, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 31: log Z = -0.5505, l_oo = {'~': {'bias_scale_param': Array(0.04734327, dtype=float32), 'flat_params_param': Array(0.02919021, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 32: log Z = -0.5497, l_oo = {'~': {'bias_scale_param': Array(0.04626544, dtype=float32), 'flat_params_param': Array(0.02953354, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 32: log Z = -0.5497, l_oo = {'~': {'bias_scale_param': Array(0.04626544, dtype=float32), 'flat_params_param': Array(0.02953354, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 33: log Z = -0.5489, l_oo = {'~': {'bias_scale_param': Array(0.04467405, dtype=float32), 'flat_params_param': Array(0.03017787, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 33: log Z = -0.5489, l_oo = {'~': {'bias_scale_param': Array(0.04467405, dtype=float32), 'flat_params_param': Array(0.03017787, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 34: log Z = -0.5482, l_oo = {'~': {'bias_scale_param': Array(0.04317542, dtype=float32), 'flat_params_param': Array(0.03076367, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 34: log Z = -0.5482, l_oo = {'~': {'bias_scale_param': Array(0.04317542, dtype=float32), 'flat_params_param': Array(0.03076367, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 35: log Z = -0.5475, l_oo = {'~': {'bias_scale_param': Array(0.0422819, dtype=float32), 'flat_params_param': Array(0.03100275, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 35: log Z = -0.5475, l_oo = {'~': {'bias_scale_param': Array(0.0422819, dtype=float32), 'flat_params_param': Array(0.03100275, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 36: log Z = -0.5467, l_oo = {'~': {'bias_scale_param': Array(0.04222003, dtype=float32), 'flat_params_param': Array(0.03077423, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 36: log Z = -0.5467, l_oo = {'~': {'bias_scale_param': Array(0.04222003, dtype=float32), 'flat_params_param': Array(0.03077423, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 37: log Z = -0.5460, l_oo = {'~': {'bias_scale_param': Array(0.04287354, dtype=float32), 'flat_params_param': Array(0.03014553, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 37: log Z = -0.5460, l_oo = {'~': {'bias_scale_param': Array(0.04287354, dtype=float32), 'flat_params_param': Array(0.03014553, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 38: log Z = -0.5452, l_oo = {'~': {'bias_scale_param': Array(0.04385115, dtype=float32), 'flat_params_param': Array(0.02934034, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 38: log Z = -0.5452, l_oo = {'~': {'bias_scale_param': Array(0.04385115, dtype=float32), 'flat_params_param': Array(0.02934034, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 39: log Z = -0.5445, l_oo = {'~': {'bias_scale_param': Array(0.04464426, dtype=float32), 'flat_params_param': Array(0.02865554, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 39: log Z = -0.5445, l_oo = {'~': {'bias_scale_param': Array(0.04464426, dtype=float32), 'flat_params_param': Array(0.02865554, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 40: log Z = -0.5438, l_oo = {'~': {'bias_scale_param': Array(0.04484976, dtype=float32), 'flat_params_param': Array(0.02834048, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 40: log Z = -0.5438, l_oo = {'~': {'bias_scale_param': Array(0.04484976, dtype=float32), 'flat_params_param': Array(0.02834048, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 41: log Z = -0.5431, l_oo = {'~': {'bias_scale_param': Array(0.04436592, dtype=float32), 'flat_params_param': Array(0.02845969, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 41: log Z = -0.5431, l_oo = {'~': {'bias_scale_param': Array(0.04436592, dtype=float32), 'flat_params_param': Array(0.02845969, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 42: log Z = -0.5424, l_oo = {'~': {'bias_scale_param': Array(0.04343135, dtype=float32), 'flat_params_param': Array(0.02886561, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 42: log Z = -0.5424, l_oo = {'~': {'bias_scale_param': Array(0.04343135, dtype=float32), 'flat_params_param': Array(0.02886561, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 43: log Z = -0.5416, l_oo = {'~': {'bias_scale_param': Array(0.04248586, dtype=float32), 'flat_params_param': Array(0.02928376, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 43: log Z = -0.5416, l_oo = {'~': {'bias_scale_param': Array(0.04248586, dtype=float32), 'flat_params_param': Array(0.02928376, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 44: log Z = -0.5409, l_oo = {'~': {'bias_scale_param': Array(0.04194589, dtype=float32), 'flat_params_param': Array(0.02946812, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 44: log Z = -0.5409, l_oo = {'~': {'bias_scale_param': Array(0.04194589, dtype=float32), 'flat_params_param': Array(0.02946812, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 45: log Z = -0.5402, l_oo = {'~': {'bias_scale_param': Array(0.04200481, dtype=float32), 'flat_params_param': Array(0.02930043, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 45: log Z = -0.5402, l_oo = {'~': {'bias_scale_param': Array(0.04200481, dtype=float32), 'flat_params_param': Array(0.02930043, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 46: log Z = -0.5395, l_oo = {'~': {'bias_scale_param': Array(0.04255427, dtype=float32), 'flat_params_param': Array(0.02883918, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 46: log Z = -0.5395, l_oo = {'~': {'bias_scale_param': Array(0.04255427, dtype=float32), 'flat_params_param': Array(0.02883918, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 47: log Z = -0.5388, l_oo = {'~': {'bias_scale_param': Array(0.04324889, dtype=float32), 'flat_params_param': Array(0.02827879, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 47: log Z = -0.5388, l_oo = {'~': {'bias_scale_param': Array(0.04324889, dtype=float32), 'flat_params_param': Array(0.02827879, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 48: log Z = -0.5381, l_oo = {'~': {'bias_scale_param': Array(0.0436766, dtype=float32), 'flat_params_param': Array(0.02786715, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 48: log Z = -0.5381, l_oo = {'~': {'bias_scale_param': Array(0.0436766, dtype=float32), 'flat_params_param': Array(0.02786715, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 49: log Z = -0.5375, l_oo = {'~': {'bias_scale_param': Array(0.04357674, dtype=float32), 'flat_params_param': Array(0.02776342, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 49: log Z = -0.5375, l_oo = {'~': {'bias_scale_param': Array(0.04357674, dtype=float32), 'flat_params_param': Array(0.02776342, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 50: log Z = -0.5368, l_oo = {'~': {'bias_scale_param': Array(0.04298979, dtype=float32), 'flat_params_param': Array(0.02794334, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 50: log Z = -0.5368, l_oo = {'~': {'bias_scale_param': Array(0.04298979, dtype=float32), 'flat_params_param': Array(0.02794334, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 51: log Z = -0.5361, l_oo = {'~': {'bias_scale_param': Array(0.04221554, dtype=float32), 'flat_params_param': Array(0.02822436, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 51: log Z = -0.5361, l_oo = {'~': {'bias_scale_param': Array(0.04221554, dtype=float32), 'flat_params_param': Array(0.02822436, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 52: log Z = -0.5354, l_oo = {'~': {'bias_scale_param': Array(0.04162725, dtype=float32), 'flat_params_param': Array(0.02837497, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 52: log Z = -0.5354, l_oo = {'~': {'bias_scale_param': Array(0.04162725, dtype=float32), 'flat_params_param': Array(0.02837497, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 53: log Z = -0.5347, l_oo = {'~': {'bias_scale_param': Array(0.04145681, dtype=float32), 'flat_params_param': Array(0.02825988, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 53: log Z = -0.5347, l_oo = {'~': {'bias_scale_param': Array(0.04145681, dtype=float32), 'flat_params_param': Array(0.02825988, dtype=float32)}}:  28%|██▊       | 28/100 [00:00<00:01, 54.36it/s]
Step 53: log Z = -0.5347, l_oo = {'~': {'bias_scale_param': Array(0.04145681, dtype=float32), 'flat_params_param': Array(0.02825988, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 54: log Z = -0.5341, l_oo = {'~': {'bias_scale_param': Array(0.04167304, dtype=float32), 'flat_params_param': Array(0.02789506, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 54: log Z = -0.5341, l_oo = {'~': {'bias_scale_param': Array(0.04167304, dtype=float32), 'flat_params_param': Array(0.02789506, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 55: log Z = -0.5334, l_oo = {'~': {'bias_scale_param': Array(0.04201547, dtype=float32), 'flat_params_param': Array(0.02743948, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 55: log Z = -0.5334, l_oo = {'~': {'bias_scale_param': Array(0.04201547, dtype=float32), 'flat_params_param': Array(0.02743948, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 56: log Z = -0.5327, l_oo = {'~': {'bias_scale_param': Array(0.04215352, dtype=float32), 'flat_params_param': Array(0.0270954, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 56: log Z = -0.5327, l_oo = {'~': {'bias_scale_param': Array(0.04215352, dtype=float32), 'flat_params_param': Array(0.0270954, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 57: log Z = -0.5321, l_oo = {'~': {'bias_scale_param': Array(0.04189419, dtype=float32), 'flat_params_param': Array(0.02698798, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 57: log Z = -0.5321, l_oo = {'~': {'bias_scale_param': Array(0.04189419, dtype=float32), 'flat_params_param': Array(0.02698798, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 58: log Z = -0.5314, l_oo = {'~': {'bias_scale_param': Array(0.04130096, dtype=float32), 'flat_params_param': Array(0.02708163, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 58: log Z = -0.5314, l_oo = {'~': {'bias_scale_param': Array(0.04130096, dtype=float32), 'flat_params_param': Array(0.02708163, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 59: log Z = -0.5308, l_oo = {'~': {'bias_scale_param': Array(0.0406374, dtype=float32), 'flat_params_param': Array(0.02721507, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 59: log Z = -0.5308, l_oo = {'~': {'bias_scale_param': Array(0.0406374, dtype=float32), 'flat_params_param': Array(0.02721507, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 60: log Z = -0.5301, l_oo = {'~': {'bias_scale_param': Array(0.04018862, dtype=float32), 'flat_params_param': Array(0.02721517, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 60: log Z = -0.5301, l_oo = {'~': {'bias_scale_param': Array(0.04018862, dtype=float32), 'flat_params_param': Array(0.02721517, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 61: log Z = -0.5295, l_oo = {'~': {'bias_scale_param': Array(0.04007779, dtype=float32), 'flat_params_param': Array(0.02700992, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 61: log Z = -0.5295, l_oo = {'~': {'bias_scale_param': Array(0.04007779, dtype=float32), 'flat_params_param': Array(0.02700992, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 62: log Z = -0.5289, l_oo = {'~': {'bias_scale_param': Array(0.04020188, dtype=float32), 'flat_params_param': Array(0.02666078, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 62: log Z = -0.5289, l_oo = {'~': {'bias_scale_param': Array(0.04020188, dtype=float32), 'flat_params_param': Array(0.02666078, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 63: log Z = -0.5282, l_oo = {'~': {'bias_scale_param': Array(0.04030971, dtype=float32), 'flat_params_param': Array(0.02632364, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 63: log Z = -0.5282, l_oo = {'~': {'bias_scale_param': Array(0.04030971, dtype=float32), 'flat_params_param': Array(0.02632364, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 64: log Z = -0.5276, l_oo = {'~': {'bias_scale_param': Array(0.04018042, dtype=float32), 'flat_params_param': Array(0.02613928, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 64: log Z = -0.5276, l_oo = {'~': {'bias_scale_param': Array(0.04018042, dtype=float32), 'flat_params_param': Array(0.02613928, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 65: log Z = -0.5270, l_oo = {'~': {'bias_scale_param': Array(0.03977285, dtype=float32), 'flat_params_param': Array(0.02613414, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 65: log Z = -0.5270, l_oo = {'~': {'bias_scale_param': Array(0.03977285, dtype=float32), 'flat_params_param': Array(0.02613414, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 66: log Z = -0.5264, l_oo = {'~': {'bias_scale_param': Array(0.03924634, dtype=float32), 'flat_params_param': Array(0.02620868, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 66: log Z = -0.5264, l_oo = {'~': {'bias_scale_param': Array(0.03924634, dtype=float32), 'flat_params_param': Array(0.02620868, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 67: log Z = -0.5258, l_oo = {'~': {'bias_scale_param': Array(0.03883454, dtype=float32), 'flat_params_param': Array(0.02621694, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 67: log Z = -0.5258, l_oo = {'~': {'bias_scale_param': Array(0.03883454, dtype=float32), 'flat_params_param': Array(0.02621694, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 68: log Z = -0.5252, l_oo = {'~': {'bias_scale_param': Array(0.03867371, dtype=float32), 'flat_params_param': Array(0.02607379, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 68: log Z = -0.5252, l_oo = {'~': {'bias_scale_param': Array(0.03867371, dtype=float32), 'flat_params_param': Array(0.02607379, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 69: log Z = -0.5246, l_oo = {'~': {'bias_scale_param': Array(0.03871207, dtype=float32), 'flat_params_param': Array(0.0258073, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 69: log Z = -0.5246, l_oo = {'~': {'bias_scale_param': Array(0.03871207, dtype=float32), 'flat_params_param': Array(0.0258073, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 70: log Z = -0.5240, l_oo = {'~': {'bias_scale_param': Array(0.03875708, dtype=float32), 'flat_params_param': Array(0.02553862, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 70: log Z = -0.5240, l_oo = {'~': {'bias_scale_param': Array(0.03875708, dtype=float32), 'flat_params_param': Array(0.02553862, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 71: log Z = -0.5234, l_oo = {'~': {'bias_scale_param': Array(0.03862428, dtype=float32), 'flat_params_param': Array(0.02538175, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 71: log Z = -0.5234, l_oo = {'~': {'bias_scale_param': Array(0.03862428, dtype=float32), 'flat_params_param': Array(0.02538175, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 72: log Z = -0.5228, l_oo = {'~': {'bias_scale_param': Array(0.03827856, dtype=float32), 'flat_params_param': Array(0.0253608, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 72: log Z = -0.5228, l_oo = {'~': {'bias_scale_param': Array(0.03827856, dtype=float32), 'flat_params_param': Array(0.0253608, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 73: log Z = -0.5222, l_oo = {'~': {'bias_scale_param': Array(0.03785126, dtype=float32), 'flat_params_param': Array(0.02539072, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 73: log Z = -0.5222, l_oo = {'~': {'bias_scale_param': Array(0.03785126, dtype=float32), 'flat_params_param': Array(0.02539072, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 74: log Z = -0.5216, l_oo = {'~': {'bias_scale_param': Array(0.03752539, dtype=float32), 'flat_params_param': Array(0.02535425, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 74: log Z = -0.5216, l_oo = {'~': {'bias_scale_param': Array(0.03752539, dtype=float32), 'flat_params_param': Array(0.02535425, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 75: log Z = -0.5211, l_oo = {'~': {'bias_scale_param': Array(0.03738682, dtype=float32), 'flat_params_param': Array(0.02519514, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 75: log Z = -0.5211, l_oo = {'~': {'bias_scale_param': Array(0.03738682, dtype=float32), 'flat_params_param': Array(0.02519514, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 76: log Z = -0.5205, l_oo = {'~': {'bias_scale_param': Array(0.03736423, dtype=float32), 'flat_params_param': Array(0.02495802, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 76: log Z = -0.5205, l_oo = {'~': {'bias_scale_param': Array(0.03736423, dtype=float32), 'flat_params_param': Array(0.02495802, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 77: log Z = -0.5199, l_oo = {'~': {'bias_scale_param': Array(0.03729281, dtype=float32), 'flat_params_param': Array(0.02474771, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 77: log Z = -0.5199, l_oo = {'~': {'bias_scale_param': Array(0.03729281, dtype=float32), 'flat_params_param': Array(0.02474771, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 78: log Z = -0.5194, l_oo = {'~': {'bias_scale_param': Array(0.0370569, dtype=float32), 'flat_params_param': Array(0.02463884, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 78: log Z = -0.5194, l_oo = {'~': {'bias_scale_param': Array(0.0370569, dtype=float32), 'flat_params_param': Array(0.02463884, dtype=float32)}}:  54%|█████▍    | 54/100 [00:00<00:00, 99.09it/s]
Step 78: log Z = -0.5194, l_oo = {'~': {'bias_scale_param': Array(0.0370569, dtype=float32), 'flat_params_param': Array(0.02463884, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 79: log Z = -0.5188, l_oo = {'~': {'bias_scale_param': Array(0.0366841, dtype=float32), 'flat_params_param': Array(0.02461465, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 79: log Z = -0.5188, l_oo = {'~': {'bias_scale_param': Array(0.0366841, dtype=float32), 'flat_params_param': Array(0.02461465, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 80: log Z = -0.5183, l_oo = {'~': {'bias_scale_param': Array(0.0363136, dtype=float32), 'flat_params_param': Array(0.02458714, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 80: log Z = -0.5183, l_oo = {'~': {'bias_scale_param': Array(0.0363136, dtype=float32), 'flat_params_param': Array(0.02458714, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 81: log Z = -0.5178, l_oo = {'~': {'bias_scale_param': Array(0.03606647, dtype=float32), 'flat_params_param': Array(0.02447585, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 81: log Z = -0.5178, l_oo = {'~': {'bias_scale_param': Array(0.03606647, dtype=float32), 'flat_params_param': Array(0.02447585, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 82: log Z = -0.5172, l_oo = {'~': {'bias_scale_param': Array(0.0359465, dtype=float32), 'flat_params_param': Array(0.02428117, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 82: log Z = -0.5172, l_oo = {'~': {'bias_scale_param': Array(0.0359465, dtype=float32), 'flat_params_param': Array(0.02428117, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 83: log Z = -0.5167, l_oo = {'~': {'bias_scale_param': Array(0.03584221, dtype=float32), 'flat_params_param': Array(0.02407176, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 83: log Z = -0.5167, l_oo = {'~': {'bias_scale_param': Array(0.03584221, dtype=float32), 'flat_params_param': Array(0.02407176, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 84: log Z = -0.5162, l_oo = {'~': {'bias_scale_param': Array(0.0356359, dtype=float32), 'flat_params_param': Array(0.02392969, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 84: log Z = -0.5162, l_oo = {'~': {'bias_scale_param': Array(0.0356359, dtype=float32), 'flat_params_param': Array(0.02392969, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 85: log Z = -0.5156, l_oo = {'~': {'bias_scale_param': Array(0.03531031, dtype=float32), 'flat_params_param': Array(0.02386588, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 85: log Z = -0.5156, l_oo = {'~': {'bias_scale_param': Array(0.03531031, dtype=float32), 'flat_params_param': Array(0.02386588, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 86: log Z = -0.5151, l_oo = {'~': {'bias_scale_param': Array(0.03495835, dtype=float32), 'flat_params_param': Array(0.02381913, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 86: log Z = -0.5151, l_oo = {'~': {'bias_scale_param': Array(0.03495835, dtype=float32), 'flat_params_param': Array(0.02381913, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 87: log Z = -0.5146, l_oo = {'~': {'bias_scale_param': Array(0.03468978, dtype=float32), 'flat_params_param': Array(0.0237208, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 87: log Z = -0.5146, l_oo = {'~': {'bias_scale_param': Array(0.03468978, dtype=float32), 'flat_params_param': Array(0.0237208, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 88: log Z = -0.5141, l_oo = {'~': {'bias_scale_param': Array(0.03452826, dtype=float32), 'flat_params_param': Array(0.02355227, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 88: log Z = -0.5141, l_oo = {'~': {'bias_scale_param': Array(0.03452826, dtype=float32), 'flat_params_param': Array(0.02355227, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 89: log Z = -0.5136, l_oo = {'~': {'bias_scale_param': Array(0.03439544, dtype=float32), 'flat_params_param': Array(0.02336685, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 89: log Z = -0.5136, l_oo = {'~': {'bias_scale_param': Array(0.03439544, dtype=float32), 'flat_params_param': Array(0.02336685, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 90: log Z = -0.5131, l_oo = {'~': {'bias_scale_param': Array(0.03419412, dtype=float32), 'flat_params_param': Array(0.023229, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 90: log Z = -0.5131, l_oo = {'~': {'bias_scale_param': Array(0.03419412, dtype=float32), 'flat_params_param': Array(0.023229, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 91: log Z = -0.5126, l_oo = {'~': {'bias_scale_param': Array(0.0339027, dtype=float32), 'flat_params_param': Array(0.0231531, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 91: log Z = -0.5126, l_oo = {'~': {'bias_scale_param': Array(0.0339027, dtype=float32), 'flat_params_param': Array(0.0231531, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 92: log Z = -0.5121, l_oo = {'~': {'bias_scale_param': Array(0.03359128, dtype=float32), 'flat_params_param': Array(0.02309234, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 92: log Z = -0.5121, l_oo = {'~': {'bias_scale_param': Array(0.03359128, dtype=float32), 'flat_params_param': Array(0.02309234, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 93: log Z = -0.5116, l_oo = {'~': {'bias_scale_param': Array(0.03334535, dtype=float32), 'flat_params_param': Array(0.02299092, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 93: log Z = -0.5116, l_oo = {'~': {'bias_scale_param': Array(0.03334535, dtype=float32), 'flat_params_param': Array(0.02299092, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 94: log Z = -0.5112, l_oo = {'~': {'bias_scale_param': Array(0.03317994, dtype=float32), 'flat_params_param': Array(0.02283694, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 94: log Z = -0.5112, l_oo = {'~': {'bias_scale_param': Array(0.03317994, dtype=float32), 'flat_params_param': Array(0.02283694, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 95: log Z = -0.5107, l_oo = {'~': {'bias_scale_param': Array(0.03302866, dtype=float32), 'flat_params_param': Array(0.02267506, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 95: log Z = -0.5107, l_oo = {'~': {'bias_scale_param': Array(0.03302866, dtype=float32), 'flat_params_param': Array(0.02267506, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 96: log Z = -0.5102, l_oo = {'~': {'bias_scale_param': Array(0.03281669, dtype=float32), 'flat_params_param': Array(0.02255439, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 96: log Z = -0.5102, l_oo = {'~': {'bias_scale_param': Array(0.03281669, dtype=float32), 'flat_params_param': Array(0.02255439, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 97: log Z = -0.5098, l_oo = {'~': {'bias_scale_param': Array(0.03254045, dtype=float32), 'flat_params_param': Array(0.02247838, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 97: log Z = -0.5098, l_oo = {'~': {'bias_scale_param': Array(0.03254045, dtype=float32), 'flat_params_param': Array(0.02247838, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 98: log Z = -0.5093, l_oo = {'~': {'bias_scale_param': Array(0.03226357, dtype=float32), 'flat_params_param': Array(0.02240312, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 98: log Z = -0.5093, l_oo = {'~': {'bias_scale_param': Array(0.03226357, dtype=float32), 'flat_params_param': Array(0.02240312, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 99: log Z = -0.5088, l_oo = {'~': {'bias_scale_param': Array(0.03204567, dtype=float32), 'flat_params_param': Array(0.02228846, dtype=float32)}}:  79%|███████▉  | 79/100 [00:00<00:00, 134.52it/s]
Step 99: log Z = -0.5088, l_oo = {'~': {'bias_scale_param': Array(0.03204567, dtype=float32), 'flat_params_param': Array(0.02228846, dtype=float32)}}: 100%|██████████| 100/100 [00:00<00:00, 102.41it/s][A
Step 4: log Z = -0.4964 +- 0.0208:  40%|████      | 4/10 [00:10<00:14,  2.37s/it]
Training Progress:   0%|          | 0/100 [00:00<?, ?it/s]
Step 0: log Z = -0.4964, l_oo = {'~': {'bias_scale_param': Array(0.01767908, dtype=float32), 'flat_params_param': Array(0.04049821, dtype=float32)}}:   0%|          | 0/100 [00:00<?, ?it/s]
Step 0: log Z = -0.4964, l_oo = {'~': {'bias_scale_param': Array(0.01767908, dtype=float32), 'flat_params_param': Array(0.04049821, dtype=float32)}}:   0%|          | 0/100 [00:00<?, ?it/s]
Step 0: log Z = -0.4964, l_oo = {'~': {'bias_scale_param': Array(0.01767908, dtype=float32), 'flat_params_param': Array(0.04049821, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 1: log Z = -0.4987, l_oo = {'~': {'bias_scale_param': Array(0.02749027, dtype=float32), 'flat_params_param': Array(0.18030968, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 1: log Z = -0.4987, l_oo = {'~': {'bias_scale_param': Array(0.02749027, dtype=float32), 'flat_params_param': Array(0.18030968, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 2: log Z = -0.4954, l_oo = {'~': {'bias_scale_param': Array(0.02243771, dtype=float32), 'flat_params_param': Array(0.07589417, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 2: log Z = -0.4954, l_oo = {'~': {'bias_scale_param': Array(0.02243771, dtype=float32), 'flat_params_param': Array(0.07589417, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 3: log Z = -0.4951, l_oo = {'~': {'bias_scale_param': Array(0.01717184, dtype=float32), 'flat_params_param': Array(0.04270048, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 3: log Z = -0.4951, l_oo = {'~': {'bias_scale_param': Array(0.01717184, dtype=float32), 'flat_params_param': Array(0.04270048, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 4: log Z = -0.4957, l_oo = {'~': {'bias_scale_param': Array(0.01505365, dtype=float32), 'flat_params_param': Array(0.08635476, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 4: log Z = -0.4957, l_oo = {'~': {'bias_scale_param': Array(0.01505365, dtype=float32), 'flat_params_param': Array(0.08635476, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 5: log Z = -0.4948, l_oo = {'~': {'bias_scale_param': Array(0.01589193, dtype=float32), 'flat_params_param': Array(0.0701948, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 5: log Z = -0.4948, l_oo = {'~': {'bias_scale_param': Array(0.01589193, dtype=float32), 'flat_params_param': Array(0.0701948, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 6: log Z = -0.4934, l_oo = {'~': {'bias_scale_param': Array(0.01837831, dtype=float32), 'flat_params_param': Array(0.02537879, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 6: log Z = -0.4934, l_oo = {'~': {'bias_scale_param': Array(0.01837831, dtype=float32), 'flat_params_param': Array(0.02537879, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 7: log Z = -0.4929, l_oo = {'~': {'bias_scale_param': Array(0.02129691, dtype=float32), 'flat_params_param': Array(0.04680825, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 7: log Z = -0.4929, l_oo = {'~': {'bias_scale_param': Array(0.02129691, dtype=float32), 'flat_params_param': Array(0.04680825, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 8: log Z = -0.4931, l_oo = {'~': {'bias_scale_param': Array(0.02325601, dtype=float32), 'flat_params_param': Array(0.08924498, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 8: log Z = -0.4931, l_oo = {'~': {'bias_scale_param': Array(0.02325601, dtype=float32), 'flat_params_param': Array(0.08924498, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 9: log Z = -0.4927, l_oo = {'~': {'bias_scale_param': Array(0.0233702, dtype=float32), 'flat_params_param': Array(0.09136917, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 9: log Z = -0.4927, l_oo = {'~': {'bias_scale_param': Array(0.0233702, dtype=float32), 'flat_params_param': Array(0.09136917, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 10: log Z = -0.4917, l_oo = {'~': {'bias_scale_param': Array(0.0219335, dtype=float32), 'flat_params_param': Array(0.05901076, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 10: log Z = -0.4917, l_oo = {'~': {'bias_scale_param': Array(0.0219335, dtype=float32), 'flat_params_param': Array(0.05901076, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 11: log Z = -0.4910, l_oo = {'~': {'bias_scale_param': Array(0.01980693, dtype=float32), 'flat_params_param': Array(0.0216732, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 11: log Z = -0.4910, l_oo = {'~': {'bias_scale_param': Array(0.01980693, dtype=float32), 'flat_params_param': Array(0.0216732, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 12: log Z = -0.4908, l_oo = {'~': {'bias_scale_param': Array(0.01791258, dtype=float32), 'flat_params_param': Array(0.03360395, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 12: log Z = -0.4908, l_oo = {'~': {'bias_scale_param': Array(0.01791258, dtype=float32), 'flat_params_param': Array(0.03360395, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 13: log Z = -0.4907, l_oo = {'~': {'bias_scale_param': Array(0.01692645, dtype=float32), 'flat_params_param': Array(0.05693386, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 13: log Z = -0.4907, l_oo = {'~': {'bias_scale_param': Array(0.01692645, dtype=float32), 'flat_params_param': Array(0.05693386, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 14: log Z = -0.4902, l_oo = {'~': {'bias_scale_param': Array(0.01704388, dtype=float32), 'flat_params_param': Array(0.05383002, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 14: log Z = -0.4902, l_oo = {'~': {'bias_scale_param': Array(0.01704388, dtype=float32), 'flat_params_param': Array(0.05383002, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 15: log Z = -0.4895, l_oo = {'~': {'bias_scale_param': Array(0.01803935, dtype=float32), 'flat_params_param': Array(0.02898403, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 15: log Z = -0.4895, l_oo = {'~': {'bias_scale_param': Array(0.01803935, dtype=float32), 'flat_params_param': Array(0.02898403, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 16: log Z = -0.4889, l_oo = {'~': {'bias_scale_param': Array(0.01943686, dtype=float32), 'flat_params_param': Array(0.02099863, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 16: log Z = -0.4889, l_oo = {'~': {'bias_scale_param': Array(0.01943686, dtype=float32), 'flat_params_param': Array(0.02099863, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 17: log Z = -0.4887, l_oo = {'~': {'bias_scale_param': Array(0.02061773, dtype=float32), 'flat_params_param': Array(0.03756502, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 17: log Z = -0.4887, l_oo = {'~': {'bias_scale_param': Array(0.02061773, dtype=float32), 'flat_params_param': Array(0.03756502, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 18: log Z = -0.4885, l_oo = {'~': {'bias_scale_param': Array(0.0210366, dtype=float32), 'flat_params_param': Array(0.05100986, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 18: log Z = -0.4885, l_oo = {'~': {'bias_scale_param': Array(0.0210366, dtype=float32), 'flat_params_param': Array(0.05100986, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 19: log Z = -0.4880, l_oo = {'~': {'bias_scale_param': Array(0.02052885, dtype=float32), 'flat_params_param': Array(0.04214492, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 19: log Z = -0.4880, l_oo = {'~': {'bias_scale_param': Array(0.02052885, dtype=float32), 'flat_params_param': Array(0.04214492, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 20: log Z = -0.4874, l_oo = {'~': {'bias_scale_param': Array(0.01936338, dtype=float32), 'flat_params_param': Array(0.01994881, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 20: log Z = -0.4874, l_oo = {'~': {'bias_scale_param': Array(0.01936338, dtype=float32), 'flat_params_param': Array(0.01994881, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 21: log Z = -0.4870, l_oo = {'~': {'bias_scale_param': Array(0.01803923, dtype=float32), 'flat_params_param': Array(0.02097611, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 21: log Z = -0.4870, l_oo = {'~': {'bias_scale_param': Array(0.01803923, dtype=float32), 'flat_params_param': Array(0.02097611, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 22: log Z = -0.4868, l_oo = {'~': {'bias_scale_param': Array(0.01705123, dtype=float32), 'flat_params_param': Array(0.03438557, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 22: log Z = -0.4868, l_oo = {'~': {'bias_scale_param': Array(0.01705123, dtype=float32), 'flat_params_param': Array(0.03438557, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 23: log Z = -0.4864, l_oo = {'~': {'bias_scale_param': Array(0.0166905, dtype=float32), 'flat_params_param': Array(0.03970561, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 23: log Z = -0.4864, l_oo = {'~': {'bias_scale_param': Array(0.0166905, dtype=float32), 'flat_params_param': Array(0.03970561, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 24: log Z = -0.4860, l_oo = {'~': {'bias_scale_param': Array(0.01696068, dtype=float32), 'flat_params_param': Array(0.02812317, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 24: log Z = -0.4860, l_oo = {'~': {'bias_scale_param': Array(0.01696068, dtype=float32), 'flat_params_param': Array(0.02812317, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 25: log Z = -0.4855, l_oo = {'~': {'bias_scale_param': Array(0.01762305, dtype=float32), 'flat_params_param': Array(0.02008336, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 25: log Z = -0.4855, l_oo = {'~': {'bias_scale_param': Array(0.01762305, dtype=float32), 'flat_params_param': Array(0.02008336, dtype=float32)}}:   1%|          | 1/100 [00:00<00:58,  1.69it/s]
Step 25: log Z = -0.4855, l_oo = {'~': {'bias_scale_param': Array(0.01762305, dtype=float32), 'flat_params_param': Array(0.02008336, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 26: log Z = -0.4852, l_oo = {'~': {'bias_scale_param': Array(0.01829429, dtype=float32), 'flat_params_param': Array(0.01903914, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 26: log Z = -0.4852, l_oo = {'~': {'bias_scale_param': Array(0.01829429, dtype=float32), 'flat_params_param': Array(0.01903914, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 27: log Z = -0.4849, l_oo = {'~': {'bias_scale_param': Array(0.0185946, dtype=float32), 'flat_params_param': Array(0.03172171, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 27: log Z = -0.4849, l_oo = {'~': {'bias_scale_param': Array(0.0185946, dtype=float32), 'flat_params_param': Array(0.03172171, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 28: log Z = -0.4845, l_oo = {'~': {'bias_scale_param': Array(0.01833463, dtype=float32), 'flat_params_param': Array(0.03038944, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 28: log Z = -0.4845, l_oo = {'~': {'bias_scale_param': Array(0.01833463, dtype=float32), 'flat_params_param': Array(0.03038944, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 29: log Z = -0.4841, l_oo = {'~': {'bias_scale_param': Array(0.01760957, dtype=float32), 'flat_params_param': Array(0.01869942, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 29: log Z = -0.4841, l_oo = {'~': {'bias_scale_param': Array(0.01760957, dtype=float32), 'flat_params_param': Array(0.01869942, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 30: log Z = -0.4837, l_oo = {'~': {'bias_scale_param': Array(0.01672033, dtype=float32), 'flat_params_param': Array(0.019286, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 30: log Z = -0.4837, l_oo = {'~': {'bias_scale_param': Array(0.01672033, dtype=float32), 'flat_params_param': Array(0.019286, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 31: log Z = -0.4834, l_oo = {'~': {'bias_scale_param': Array(0.01600794, dtype=float32), 'flat_params_param': Array(0.01968666, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 31: log Z = -0.4834, l_oo = {'~': {'bias_scale_param': Array(0.01600794, dtype=float32), 'flat_params_param': Array(0.01968666, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 32: log Z = -0.4831, l_oo = {'~': {'bias_scale_param': Array(0.0156936, dtype=float32), 'flat_params_param': Array(0.02278243, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 32: log Z = -0.4831, l_oo = {'~': {'bias_scale_param': Array(0.0156936, dtype=float32), 'flat_params_param': Array(0.02278243, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 33: log Z = -0.4828, l_oo = {'~': {'bias_scale_param': Array(0.01579805, dtype=float32), 'flat_params_param': Array(0.01924429, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 33: log Z = -0.4828, l_oo = {'~': {'bias_scale_param': Array(0.01579805, dtype=float32), 'flat_params_param': Array(0.01924429, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 34: log Z = -0.4824, l_oo = {'~': {'bias_scale_param': Array(0.01615558, dtype=float32), 'flat_params_param': Array(0.01854196, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 34: log Z = -0.4824, l_oo = {'~': {'bias_scale_param': Array(0.01615558, dtype=float32), 'flat_params_param': Array(0.01854196, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 35: log Z = -0.4821, l_oo = {'~': {'bias_scale_param': Array(0.01649604, dtype=float32), 'flat_params_param': Array(0.01785317, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 35: log Z = -0.4821, l_oo = {'~': {'bias_scale_param': Array(0.01649604, dtype=float32), 'flat_params_param': Array(0.01785317, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 36: log Z = -0.4818, l_oo = {'~': {'bias_scale_param': Array(0.01657102, dtype=float32), 'flat_params_param': Array(0.02464323, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 36: log Z = -0.4818, l_oo = {'~': {'bias_scale_param': Array(0.01657102, dtype=float32), 'flat_params_param': Array(0.02464323, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 37: log Z = -0.4815, l_oo = {'~': {'bias_scale_param': Array(0.01628565, dtype=float32), 'flat_params_param': Array(0.02106229, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 37: log Z = -0.4815, l_oo = {'~': {'bias_scale_param': Array(0.01628565, dtype=float32), 'flat_params_param': Array(0.02106229, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 38: log Z = -0.4811, l_oo = {'~': {'bias_scale_param': Array(0.01574219, dtype=float32), 'flat_params_param': Array(0.01773632, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 38: log Z = -0.4811, l_oo = {'~': {'bias_scale_param': Array(0.01574219, dtype=float32), 'flat_params_param': Array(0.01773632, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 39: log Z = -0.4808, l_oo = {'~': {'bias_scale_param': Array(0.01516687, dtype=float32), 'flat_params_param': Array(0.01806469, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 39: log Z = -0.4808, l_oo = {'~': {'bias_scale_param': Array(0.01516687, dtype=float32), 'flat_params_param': Array(0.01806469, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 40: log Z = -0.4805, l_oo = {'~': {'bias_scale_param': Array(0.01478016, dtype=float32), 'flat_params_param': Array(0.01819629, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 40: log Z = -0.4805, l_oo = {'~': {'bias_scale_param': Array(0.01478016, dtype=float32), 'flat_params_param': Array(0.01819629, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 41: log Z = -0.4802, l_oo = {'~': {'bias_scale_param': Array(0.01468569, dtype=float32), 'flat_params_param': Array(0.01801915, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 41: log Z = -0.4802, l_oo = {'~': {'bias_scale_param': Array(0.01468569, dtype=float32), 'flat_params_param': Array(0.01801915, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 42: log Z = -0.4799, l_oo = {'~': {'bias_scale_param': Array(0.0148315, dtype=float32), 'flat_params_param': Array(0.0175855, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 42: log Z = -0.4799, l_oo = {'~': {'bias_scale_param': Array(0.0148315, dtype=float32), 'flat_params_param': Array(0.0175855, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 43: log Z = -0.4796, l_oo = {'~': {'bias_scale_param': Array(0.01504723, dtype=float32), 'flat_params_param': Array(0.01707615, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 43: log Z = -0.4796, l_oo = {'~': {'bias_scale_param': Array(0.01504723, dtype=float32), 'flat_params_param': Array(0.01707615, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 44: log Z = -0.4793, l_oo = {'~': {'bias_scale_param': Array(0.01513656, dtype=float32), 'flat_params_param': Array(0.016769, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 44: log Z = -0.4793, l_oo = {'~': {'bias_scale_param': Array(0.01513656, dtype=float32), 'flat_params_param': Array(0.016769, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 45: log Z = -0.4790, l_oo = {'~': {'bias_scale_param': Array(0.01498746, dtype=float32), 'flat_params_param': Array(0.01660758, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 45: log Z = -0.4790, l_oo = {'~': {'bias_scale_param': Array(0.01498746, dtype=float32), 'flat_params_param': Array(0.01660758, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 46: log Z = -0.4787, l_oo = {'~': {'bias_scale_param': Array(0.01463417, dtype=float32), 'flat_params_param': Array(0.01674188, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 46: log Z = -0.4787, l_oo = {'~': {'bias_scale_param': Array(0.01463417, dtype=float32), 'flat_params_param': Array(0.01674188, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 47: log Z = -0.4785, l_oo = {'~': {'bias_scale_param': Array(0.0142253, dtype=float32), 'flat_params_param': Array(0.01694237, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 47: log Z = -0.4785, l_oo = {'~': {'bias_scale_param': Array(0.0142253, dtype=float32), 'flat_params_param': Array(0.01694237, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 48: log Z = -0.4782, l_oo = {'~': {'bias_scale_param': Array(0.01392721, dtype=float32), 'flat_params_param': Array(0.01702148, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 48: log Z = -0.4782, l_oo = {'~': {'bias_scale_param': Array(0.01392721, dtype=float32), 'flat_params_param': Array(0.01702148, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 49: log Z = -0.4779, l_oo = {'~': {'bias_scale_param': Array(0.01382793, dtype=float32), 'flat_params_param': Array(0.01688182, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 49: log Z = -0.4779, l_oo = {'~': {'bias_scale_param': Array(0.01382793, dtype=float32), 'flat_params_param': Array(0.01688182, dtype=float32)}}:  26%|██▌       | 26/100 [00:00<00:01, 49.70it/s]
Step 49: log Z = -0.4779, l_oo = {'~': {'bias_scale_param': Array(0.01382793, dtype=float32), 'flat_params_param': Array(0.01688182, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 50: log Z = -0.4776, l_oo = {'~': {'bias_scale_param': Array(0.01389438, dtype=float32), 'flat_params_param': Array(0.01655631, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 50: log Z = -0.4776, l_oo = {'~': {'bias_scale_param': Array(0.01389438, dtype=float32), 'flat_params_param': Array(0.01655631, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 51: log Z = -0.4774, l_oo = {'~': {'bias_scale_param': Array(0.01400085, dtype=float32), 'flat_params_param': Array(0.01618329, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 51: log Z = -0.4774, l_oo = {'~': {'bias_scale_param': Array(0.01400085, dtype=float32), 'flat_params_param': Array(0.01618329, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 52: log Z = -0.4771, l_oo = {'~': {'bias_scale_param': Array(0.014006, dtype=float32), 'flat_params_param': Array(0.01592446, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 52: log Z = -0.4771, l_oo = {'~': {'bias_scale_param': Array(0.014006, dtype=float32), 'flat_params_param': Array(0.01592446, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 53: log Z = -0.4769, l_oo = {'~': {'bias_scale_param': Array(0.01384197, dtype=float32), 'flat_params_param': Array(0.01586181, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 53: log Z = -0.4769, l_oo = {'~': {'bias_scale_param': Array(0.01384197, dtype=float32), 'flat_params_param': Array(0.01586181, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 54: log Z = -0.4766, l_oo = {'~': {'bias_scale_param': Array(0.01355221, dtype=float32), 'flat_params_param': Array(0.01594332, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 54: log Z = -0.4766, l_oo = {'~': {'bias_scale_param': Array(0.01355221, dtype=float32), 'flat_params_param': Array(0.01594332, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 55: log Z = -0.4764, l_oo = {'~': {'bias_scale_param': Array(0.01325506, dtype=float32), 'flat_params_param': Array(0.01603531, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 55: log Z = -0.4764, l_oo = {'~': {'bias_scale_param': Array(0.01325506, dtype=float32), 'flat_params_param': Array(0.01603531, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 56: log Z = -0.4761, l_oo = {'~': {'bias_scale_param': Array(0.01306216, dtype=float32), 'flat_params_param': Array(0.0160082, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 56: log Z = -0.4761, l_oo = {'~': {'bias_scale_param': Array(0.01306216, dtype=float32), 'flat_params_param': Array(0.0160082, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 57: log Z = -0.4759, l_oo = {'~': {'bias_scale_param': Array(0.01300884, dtype=float32), 'flat_params_param': Array(0.01582062, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 57: log Z = -0.4759, l_oo = {'~': {'bias_scale_param': Array(0.01300884, dtype=float32), 'flat_params_param': Array(0.01582062, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 58: log Z = -0.4756, l_oo = {'~': {'bias_scale_param': Array(0.01303908, dtype=float32), 'flat_params_param': Array(0.01553611, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 58: log Z = -0.4756, l_oo = {'~': {'bias_scale_param': Array(0.01303908, dtype=float32), 'flat_params_param': Array(0.01553611, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 59: log Z = -0.4754, l_oo = {'~': {'bias_scale_param': Array(0.013048, dtype=float32), 'flat_params_param': Array(0.01527502, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 59: log Z = -0.4754, l_oo = {'~': {'bias_scale_param': Array(0.013048, dtype=float32), 'flat_params_param': Array(0.01527502, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 60: log Z = -0.4751, l_oo = {'~': {'bias_scale_param': Array(0.01295277, dtype=float32), 'flat_params_param': Array(0.01513617, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 60: log Z = -0.4751, l_oo = {'~': {'bias_scale_param': Array(0.01295277, dtype=float32), 'flat_params_param': Array(0.01513617, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 61: log Z = -0.4749, l_oo = {'~': {'bias_scale_param': Array(0.01274829, dtype=float32), 'flat_params_param': Array(0.01512794, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 61: log Z = -0.4749, l_oo = {'~': {'bias_scale_param': Array(0.01274829, dtype=float32), 'flat_params_param': Array(0.01512794, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 62: log Z = -0.4747, l_oo = {'~': {'bias_scale_param': Array(0.01250396, dtype=float32), 'flat_params_param': Array(0.01516819, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 62: log Z = -0.4747, l_oo = {'~': {'bias_scale_param': Array(0.01250396, dtype=float32), 'flat_params_param': Array(0.01516819, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 63: log Z = -0.4744, l_oo = {'~': {'bias_scale_param': Array(0.0123104, dtype=float32), 'flat_params_param': Array(0.01515046, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 63: log Z = -0.4744, l_oo = {'~': {'bias_scale_param': Array(0.0123104, dtype=float32), 'flat_params_param': Array(0.01515046, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 64: log Z = -0.4742, l_oo = {'~': {'bias_scale_param': Array(0.01221458, dtype=float32), 'flat_params_param': Array(0.0150203, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 64: log Z = -0.4742, l_oo = {'~': {'bias_scale_param': Array(0.01221458, dtype=float32), 'flat_params_param': Array(0.0150203, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 65: log Z = -0.4740, l_oo = {'~': {'bias_scale_param': Array(0.01219369, dtype=float32), 'flat_params_param': Array(0.01480212, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 65: log Z = -0.4740, l_oo = {'~': {'bias_scale_param': Array(0.01219369, dtype=float32), 'flat_params_param': Array(0.01480212, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 66: log Z = -0.4738, l_oo = {'~': {'bias_scale_param': Array(0.01217513, dtype=float32), 'flat_params_param': Array(0.01458311, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 66: log Z = -0.4738, l_oo = {'~': {'bias_scale_param': Array(0.01217513, dtype=float32), 'flat_params_param': Array(0.01458311, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 67: log Z = -0.4736, l_oo = {'~': {'bias_scale_param': Array(0.01209236, dtype=float32), 'flat_params_param': Array(0.01444285, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 67: log Z = -0.4736, l_oo = {'~': {'bias_scale_param': Array(0.01209236, dtype=float32), 'flat_params_param': Array(0.01444285, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 68: log Z = -0.4733, l_oo = {'~': {'bias_scale_param': Array(0.01193083, dtype=float32), 'flat_params_param': Array(0.01440054, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 68: log Z = -0.4733, l_oo = {'~': {'bias_scale_param': Array(0.01193083, dtype=float32), 'flat_params_param': Array(0.01440054, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 69: log Z = -0.4731, l_oo = {'~': {'bias_scale_param': Array(0.01173623, dtype=float32), 'flat_params_param': Array(0.01440185, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 69: log Z = -0.4731, l_oo = {'~': {'bias_scale_param': Array(0.01173623, dtype=float32), 'flat_params_param': Array(0.01440185, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 70: log Z = -0.4729, l_oo = {'~': {'bias_scale_param': Array(0.01157435, dtype=float32), 'flat_params_param': Array(0.01436687, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 70: log Z = -0.4729, l_oo = {'~': {'bias_scale_param': Array(0.01157435, dtype=float32), 'flat_params_param': Array(0.01436687, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 71: log Z = -0.4727, l_oo = {'~': {'bias_scale_param': Array(0.01148184, dtype=float32), 'flat_params_param': Array(0.01425105, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 71: log Z = -0.4727, l_oo = {'~': {'bias_scale_param': Array(0.01148184, dtype=float32), 'flat_params_param': Array(0.01425105, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 72: log Z = -0.4725, l_oo = {'~': {'bias_scale_param': Array(0.01144262, dtype=float32), 'flat_params_param': Array(0.01407333, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 72: log Z = -0.4725, l_oo = {'~': {'bias_scale_param': Array(0.01144262, dtype=float32), 'flat_params_param': Array(0.01407333, dtype=float32)}}:  50%|█████     | 50/100 [00:00<00:00, 90.52it/s]
Step 72: log Z = -0.4725, l_oo = {'~': {'bias_scale_param': Array(0.01144262, dtype=float32), 'flat_params_param': Array(0.01407333, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 73: log Z = -0.4723, l_oo = {'~': {'bias_scale_param': Array(0.0114032, dtype=float32), 'flat_params_param': Array(0.01389808, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 73: log Z = -0.4723, l_oo = {'~': {'bias_scale_param': Array(0.0114032, dtype=float32), 'flat_params_param': Array(0.01389808, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 74: log Z = -0.4721, l_oo = {'~': {'bias_scale_param': Array(0.01131644, dtype=float32), 'flat_params_param': Array(0.01378421, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 74: log Z = -0.4721, l_oo = {'~': {'bias_scale_param': Array(0.01131644, dtype=float32), 'flat_params_param': Array(0.01378421, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 75: log Z = -0.4719, l_oo = {'~': {'bias_scale_param': Array(0.01117677, dtype=float32), 'flat_params_param': Array(0.01373892, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 75: log Z = -0.4719, l_oo = {'~': {'bias_scale_param': Array(0.01117677, dtype=float32), 'flat_params_param': Array(0.01373892, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 76: log Z = -0.4717, l_oo = {'~': {'bias_scale_param': Array(0.01102079, dtype=float32), 'flat_params_param': Array(0.01371675, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 76: log Z = -0.4717, l_oo = {'~': {'bias_scale_param': Array(0.01102079, dtype=float32), 'flat_params_param': Array(0.01371675, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 77: log Z = -0.4715, l_oo = {'~': {'bias_scale_param': Array(0.01089583, dtype=float32), 'flat_params_param': Array(0.01365918, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 77: log Z = -0.4715, l_oo = {'~': {'bias_scale_param': Array(0.01089583, dtype=float32), 'flat_params_param': Array(0.01365918, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 78: log Z = -0.4713, l_oo = {'~': {'bias_scale_param': Array(0.01082131, dtype=float32), 'flat_params_param': Array(0.01354255, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 78: log Z = -0.4713, l_oo = {'~': {'bias_scale_param': Array(0.01082131, dtype=float32), 'flat_params_param': Array(0.01354255, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 79: log Z = -0.4712, l_oo = {'~': {'bias_scale_param': Array(0.01077644, dtype=float32), 'flat_params_param': Array(0.01339108, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 79: log Z = -0.4712, l_oo = {'~': {'bias_scale_param': Array(0.01077644, dtype=float32), 'flat_params_param': Array(0.01339108, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 80: log Z = -0.4710, l_oo = {'~': {'bias_scale_param': Array(0.01071987, dtype=float32), 'flat_params_param': Array(0.01325671, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 80: log Z = -0.4710, l_oo = {'~': {'bias_scale_param': Array(0.01071987, dtype=float32), 'flat_params_param': Array(0.01325671, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 81: log Z = -0.4708, l_oo = {'~': {'bias_scale_param': Array(0.0106244, dtype=float32), 'flat_params_param': Array(0.0131728, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 81: log Z = -0.4708, l_oo = {'~': {'bias_scale_param': Array(0.0106244, dtype=float32), 'flat_params_param': Array(0.0131728, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 82: log Z = -0.4706, l_oo = {'~': {'bias_scale_param': Array(0.01049804, dtype=float32), 'flat_params_param': Array(0.0131307, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 82: log Z = -0.4706, l_oo = {'~': {'bias_scale_param': Array(0.01049804, dtype=float32), 'flat_params_param': Array(0.0131307, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 83: log Z = -0.4704, l_oo = {'~': {'bias_scale_param': Array(0.01037397, dtype=float32), 'flat_params_param': Array(0.01308731, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 83: log Z = -0.4704, l_oo = {'~': {'bias_scale_param': Array(0.01037397, dtype=float32), 'flat_params_param': Array(0.01308731, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 84: log Z = -0.4703, l_oo = {'~': {'bias_scale_param': Array(0.01028091, dtype=float32), 'flat_params_param': Array(0.01300816, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 84: log Z = -0.4703, l_oo = {'~': {'bias_scale_param': Array(0.01028091, dtype=float32), 'flat_params_param': Array(0.01300816, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 85: log Z = -0.4701, l_oo = {'~': {'bias_scale_param': Array(0.01022029, dtype=float32), 'flat_params_param': Array(0.01288847, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 85: log Z = -0.4701, l_oo = {'~': {'bias_scale_param': Array(0.01022029, dtype=float32), 'flat_params_param': Array(0.01288847, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 86: log Z = -0.4699, l_oo = {'~': {'bias_scale_param': Array(0.01016723, dtype=float32), 'flat_params_param': Array(0.01276148, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 86: log Z = -0.4699, l_oo = {'~': {'bias_scale_param': Array(0.01016723, dtype=float32), 'flat_params_param': Array(0.01276148, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 87: log Z = -0.4697, l_oo = {'~': {'bias_scale_param': Array(0.01009342, dtype=float32), 'flat_params_param': Array(0.01266253, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 87: log Z = -0.4697, l_oo = {'~': {'bias_scale_param': Array(0.01009342, dtype=float32), 'flat_params_param': Array(0.01266253, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 88: log Z = -0.4696, l_oo = {'~': {'bias_scale_param': Array(0.00999176, dtype=float32), 'flat_params_param': Array(0.01260054, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 88: log Z = -0.4696, l_oo = {'~': {'bias_scale_param': Array(0.00999176, dtype=float32), 'flat_params_param': Array(0.01260054, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 89: log Z = -0.4694, l_oo = {'~': {'bias_scale_param': Array(0.00987995, dtype=float32), 'flat_params_param': Array(0.01255425, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 89: log Z = -0.4694, l_oo = {'~': {'bias_scale_param': Array(0.00987995, dtype=float32), 'flat_params_param': Array(0.01255425, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 90: log Z = -0.4692, l_oo = {'~': {'bias_scale_param': Array(0.00978402, dtype=float32), 'flat_params_param': Array(0.01248851, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 90: log Z = -0.4692, l_oo = {'~': {'bias_scale_param': Array(0.00978402, dtype=float32), 'flat_params_param': Array(0.01248851, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 91: log Z = -0.4691, l_oo = {'~': {'bias_scale_param': Array(0.00971373, dtype=float32), 'flat_params_param': Array(0.01239156, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 91: log Z = -0.4691, l_oo = {'~': {'bias_scale_param': Array(0.00971373, dtype=float32), 'flat_params_param': Array(0.01239156, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 92: log Z = -0.4689, l_oo = {'~': {'bias_scale_param': Array(0.00965683, dtype=float32), 'flat_params_param': Array(0.01227899, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 92: log Z = -0.4689, l_oo = {'~': {'bias_scale_param': Array(0.00965683, dtype=float32), 'flat_params_param': Array(0.01227899, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 93: log Z = -0.4688, l_oo = {'~': {'bias_scale_param': Array(0.00959084, dtype=float32), 'flat_params_param': Array(0.01217994, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 93: log Z = -0.4688, l_oo = {'~': {'bias_scale_param': Array(0.00959084, dtype=float32), 'flat_params_param': Array(0.01217994, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 94: log Z = -0.4686, l_oo = {'~': {'bias_scale_param': Array(0.00950446, dtype=float32), 'flat_params_param': Array(0.01210919, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 94: log Z = -0.4686, l_oo = {'~': {'bias_scale_param': Array(0.00950446, dtype=float32), 'flat_params_param': Array(0.01210919, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 95: log Z = -0.4685, l_oo = {'~': {'bias_scale_param': Array(0.00940605, dtype=float32), 'flat_params_param': Array(0.01205552, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 95: log Z = -0.4685, l_oo = {'~': {'bias_scale_param': Array(0.00940605, dtype=float32), 'flat_params_param': Array(0.01205552, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 96: log Z = -0.4683, l_oo = {'~': {'bias_scale_param': Array(0.00931523, dtype=float32), 'flat_params_param': Array(0.01199426, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 96: log Z = -0.4683, l_oo = {'~': {'bias_scale_param': Array(0.00931523, dtype=float32), 'flat_params_param': Array(0.01199426, dtype=float32)}}:  73%|███████▎  | 73/100 [00:00<00:00, 123.01it/s]
Step 96: log Z = -0.4683, l_oo = {'~': {'bias_scale_param': Array(0.00931523, dtype=float32), 'flat_params_param': Array(0.01199426, dtype=float32)}}:  97%|█████████▋| 97/100 [00:01<00:00, 151.39it/s]
Step 97: log Z = -0.4681, l_oo = {'~': {'bias_scale_param': Array(0.00924332, dtype=float32), 'flat_params_param': Array(0.01190993, dtype=float32)}}:  97%|█████████▋| 97/100 [00:01<00:00, 151.39it/s]
Step 97: log Z = -0.4681, l_oo = {'~': {'bias_scale_param': Array(0.00924332, dtype=float32), 'flat_params_param': Array(0.01190993, dtype=float32)}}:  97%|█████████▋| 97/100 [00:01<00:00, 151.39it/s]
Step 98: log Z = -0.4680, l_oo = {'~': {'bias_scale_param': Array(0.00918392, dtype=float32), 'flat_params_param': Array(0.01181082, dtype=float32)}}:  97%|█████████▋| 97/100 [00:01<00:00, 151.39it/s]
Step 98: log Z = -0.4680, l_oo = {'~': {'bias_scale_param': Array(0.00918392, dtype=float32), 'flat_params_param': Array(0.01181082, dtype=float32)}}:  97%|█████████▋| 97/100 [00:01<00:00, 151.39it/s]
Step 99: log Z = -0.4679, l_oo = {'~': {'bias_scale_param': Array(0.009121, dtype=float32), 'flat_params_param': Array(0.01171845, dtype=float32)}}:  97%|█████████▋| 97/100 [00:01<00:00, 151.39it/s]
Step 99: log Z = -0.4679, l_oo = {'~': {'bias_scale_param': Array(0.009121, dtype=float32), 'flat_params_param': Array(0.01171845, dtype=float32)}}: 100%|██████████| 100/100 [00:01<00:00, 98.51it/s]
Step 5: log Z = -0.4716 +- 0.0183:  50%|█████     | 5/10 [00:12<00:10,  2.08s/it]
Training Progress:   0%|          | 0/100 [00:00<?, ?it/s]
Step 0: log Z = -0.4716, l_oo = {'~': {'bias_scale_param': Array(0.0210746, dtype=float32), 'flat_params_param': Array(0.07744487, dtype=float32)}}:   0%|          | 0/100 [00:00<?, ?it/s]
Step 0: log Z = -0.4716, l_oo = {'~': {'bias_scale_param': Array(0.0210746, dtype=float32), 'flat_params_param': Array(0.07744487, dtype=float32)}}:   0%|          | 0/100 [00:00<?, ?it/s]
Step 0: log Z = -0.4716, l_oo = {'~': {'bias_scale_param': Array(0.0210746, dtype=float32), 'flat_params_param': Array(0.07744487, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 1: log Z = -0.4730, l_oo = {'~': {'bias_scale_param': Array(0.01203398, dtype=float32), 'flat_params_param': Array(0.16612343, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 1: log Z = -0.4730, l_oo = {'~': {'bias_scale_param': Array(0.01203398, dtype=float32), 'flat_params_param': Array(0.16612343, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 2: log Z = -0.4710, l_oo = {'~': {'bias_scale_param': Array(0.01361555, dtype=float32), 'flat_params_param': Array(0.11084662, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 2: log Z = -0.4710, l_oo = {'~': {'bias_scale_param': Array(0.01361555, dtype=float32), 'flat_params_param': Array(0.11084662, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 3: log Z = -0.4695, l_oo = {'~': {'bias_scale_param': Array(0.01712816, dtype=float32), 'flat_params_param': Array(0.01262927, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 3: log Z = -0.4695, l_oo = {'~': {'bias_scale_param': Array(0.01712816, dtype=float32), 'flat_params_param': Array(0.01262927, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 4: log Z = -0.4704, l_oo = {'~': {'bias_scale_param': Array(0.01952353, dtype=float32), 'flat_params_param': Array(0.08273885, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 4: log Z = -0.4704, l_oo = {'~': {'bias_scale_param': Array(0.01952353, dtype=float32), 'flat_params_param': Array(0.08273885, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 5: log Z = -0.4701, l_oo = {'~': {'bias_scale_param': Array(0.01911581, dtype=float32), 'flat_params_param': Array(0.08665251, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 5: log Z = -0.4701, l_oo = {'~': {'bias_scale_param': Array(0.01911581, dtype=float32), 'flat_params_param': Array(0.08665251, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 6: log Z = -0.4689, l_oo = {'~': {'bias_scale_param': Array(0.01695679, dtype=float32), 'flat_params_param': Array(0.03767197, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 6: log Z = -0.4689, l_oo = {'~': {'bias_scale_param': Array(0.01695679, dtype=float32), 'flat_params_param': Array(0.03767197, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 7: log Z = -0.4685, l_oo = {'~': {'bias_scale_param': Array(0.01440685, dtype=float32), 'flat_params_param': Array(0.02477777, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 7: log Z = -0.4685, l_oo = {'~': {'bias_scale_param': Array(0.01440685, dtype=float32), 'flat_params_param': Array(0.02477777, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 8: log Z = -0.4688, l_oo = {'~': {'bias_scale_param': Array(0.01258875, dtype=float32), 'flat_params_param': Array(0.06744502, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 8: log Z = -0.4688, l_oo = {'~': {'bias_scale_param': Array(0.01258875, dtype=float32), 'flat_params_param': Array(0.06744502, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 9: log Z = -0.4687, l_oo = {'~': {'bias_scale_param': Array(0.01202414, dtype=float32), 'flat_params_param': Array(0.07380445, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 9: log Z = -0.4687, l_oo = {'~': {'bias_scale_param': Array(0.01202414, dtype=float32), 'flat_params_param': Array(0.07380445, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 10: log Z = -0.4680, l_oo = {'~': {'bias_scale_param': Array(0.01254387, dtype=float32), 'flat_params_param': Array(0.04748639, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 10: log Z = -0.4680, l_oo = {'~': {'bias_scale_param': Array(0.01254387, dtype=float32), 'flat_params_param': Array(0.04748639, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 11: log Z = -0.4675, l_oo = {'~': {'bias_scale_param': Array(0.01365849, dtype=float32), 'flat_params_param': Array(0.0100116, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 11: log Z = -0.4675, l_oo = {'~': {'bias_scale_param': Array(0.01365849, dtype=float32), 'flat_params_param': Array(0.0100116, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 12: log Z = -0.4674, l_oo = {'~': {'bias_scale_param': Array(0.01475542, dtype=float32), 'flat_params_param': Array(0.04255109, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 12: log Z = -0.4674, l_oo = {'~': {'bias_scale_param': Array(0.01475542, dtype=float32), 'flat_params_param': Array(0.04255109, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 13: log Z = -0.4675, l_oo = {'~': {'bias_scale_param': Array(0.01525418, dtype=float32), 'flat_params_param': Array(0.06834789, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 13: log Z = -0.4675, l_oo = {'~': {'bias_scale_param': Array(0.01525418, dtype=float32), 'flat_params_param': Array(0.06834789, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 14: log Z = -0.4672, l_oo = {'~': {'bias_scale_param': Array(0.01490677, dtype=float32), 'flat_params_param': Array(0.06585477, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 14: log Z = -0.4672, l_oo = {'~': {'bias_scale_param': Array(0.01490677, dtype=float32), 'flat_params_param': Array(0.06585477, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 15: log Z = -0.4667, l_oo = {'~': {'bias_scale_param': Array(0.01389676, dtype=float32), 'flat_params_param': Array(0.04028415, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 15: log Z = -0.4667, l_oo = {'~': {'bias_scale_param': Array(0.01389676, dtype=float32), 'flat_params_param': Array(0.04028415, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 16: log Z = -0.4664, l_oo = {'~': {'bias_scale_param': Array(0.0126498, dtype=float32), 'flat_params_param': Array(0.00948136, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 16: log Z = -0.4664, l_oo = {'~': {'bias_scale_param': Array(0.0126498, dtype=float32), 'flat_params_param': Array(0.00948136, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 17: log Z = -0.4663, l_oo = {'~': {'bias_scale_param': Array(0.01162538, dtype=float32), 'flat_params_param': Array(0.0281523, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 17: log Z = -0.4663, l_oo = {'~': {'bias_scale_param': Array(0.01162538, dtype=float32), 'flat_params_param': Array(0.0281523, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 18: log Z = -0.4662, l_oo = {'~': {'bias_scale_param': Array(0.01113591, dtype=float32), 'flat_params_param': Array(0.0366837, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 18: log Z = -0.4662, l_oo = {'~': {'bias_scale_param': Array(0.01113591, dtype=float32), 'flat_params_param': Array(0.0366837, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 19: log Z = -0.4659, l_oo = {'~': {'bias_scale_param': Array(0.01124149, dtype=float32), 'flat_params_param': Array(0.03064726, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 19: log Z = -0.4659, l_oo = {'~': {'bias_scale_param': Array(0.01124149, dtype=float32), 'flat_params_param': Array(0.03064726, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 20: log Z = -0.4656, l_oo = {'~': {'bias_scale_param': Array(0.01177879, dtype=float32), 'flat_params_param': Array(0.01340609, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 20: log Z = -0.4656, l_oo = {'~': {'bias_scale_param': Array(0.01177879, dtype=float32), 'flat_params_param': Array(0.01340609, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 21: log Z = -0.4653, l_oo = {'~': {'bias_scale_param': Array(0.0124426, dtype=float32), 'flat_params_param': Array(0.02342744, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 21: log Z = -0.4653, l_oo = {'~': {'bias_scale_param': Array(0.0124426, dtype=float32), 'flat_params_param': Array(0.02342744, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 22: log Z = -0.4652, l_oo = {'~': {'bias_scale_param': Array(0.01288716, dtype=float32), 'flat_params_param': Array(0.04324423, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 22: log Z = -0.4652, l_oo = {'~': {'bias_scale_param': Array(0.01288716, dtype=float32), 'flat_params_param': Array(0.04324423, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 23: log Z = -0.4651, l_oo = {'~': {'bias_scale_param': Array(0.01287822, dtype=float32), 'flat_params_param': Array(0.04664754, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 23: log Z = -0.4651, l_oo = {'~': {'bias_scale_param': Array(0.01287822, dtype=float32), 'flat_params_param': Array(0.04664754, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 24: log Z = -0.4648, l_oo = {'~': {'bias_scale_param': Array(0.01241143, dtype=float32), 'flat_params_param': Array(0.03307218, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 24: log Z = -0.4648, l_oo = {'~': {'bias_scale_param': Array(0.01241143, dtype=float32), 'flat_params_param': Array(0.03307218, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 25: log Z = -0.4645, l_oo = {'~': {'bias_scale_param': Array(0.01169234, dtype=float32), 'flat_params_param': Array(0.00973414, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 25: log Z = -0.4645, l_oo = {'~': {'bias_scale_param': Array(0.01169234, dtype=float32), 'flat_params_param': Array(0.00973414, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 26: log Z = -0.4643, l_oo = {'~': {'bias_scale_param': Array(0.01101799, dtype=float32), 'flat_params_param': Array(0.01548311, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 26: log Z = -0.4643, l_oo = {'~': {'bias_scale_param': Array(0.01101799, dtype=float32), 'flat_params_param': Array(0.01548311, dtype=float32)}}:   1%|          | 1/100 [00:00<01:09,  1.43it/s]
Step 26: log Z = -0.4643, l_oo = {'~': {'bias_scale_param': Array(0.01101799, dtype=float32), 'flat_params_param': Array(0.01548311, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 27: log Z = -0.4642, l_oo = {'~': {'bias_scale_param': Array(0.01063498, dtype=float32), 'flat_params_param': Array(0.0242474, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 27: log Z = -0.4642, l_oo = {'~': {'bias_scale_param': Array(0.01063498, dtype=float32), 'flat_params_param': Array(0.0242474, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 28: log Z = -0.4640, l_oo = {'~': {'bias_scale_param': Array(0.0106372, dtype=float32), 'flat_params_param': Array(0.02159748, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 28: log Z = -0.4640, l_oo = {'~': {'bias_scale_param': Array(0.0106372, dtype=float32), 'flat_params_param': Array(0.02159748, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 29: log Z = -0.4637, l_oo = {'~': {'bias_scale_param': Array(0.01094291, dtype=float32), 'flat_params_param': Array(0.00887471, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 29: log Z = -0.4637, l_oo = {'~': {'bias_scale_param': Array(0.01094291, dtype=float32), 'flat_params_param': Array(0.00887471, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 30: log Z = -0.4636, l_oo = {'~': {'bias_scale_param': Array(0.01134114, dtype=float32), 'flat_params_param': Array(0.01111935, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 30: log Z = -0.4636, l_oo = {'~': {'bias_scale_param': Array(0.01134114, dtype=float32), 'flat_params_param': Array(0.01111935, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 31: log Z = -0.4634, l_oo = {'~': {'bias_scale_param': Array(0.01158556, dtype=float32), 'flat_params_param': Array(0.02392773, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 31: log Z = -0.4634, l_oo = {'~': {'bias_scale_param': Array(0.01158556, dtype=float32), 'flat_params_param': Array(0.02392773, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 32: log Z = -0.4633, l_oo = {'~': {'bias_scale_param': Array(0.01151561, dtype=float32), 'flat_params_param': Array(0.024735, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 32: log Z = -0.4633, l_oo = {'~': {'bias_scale_param': Array(0.01151561, dtype=float32), 'flat_params_param': Array(0.024735, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 33: log Z = -0.4630, l_oo = {'~': {'bias_scale_param': Array(0.01114203, dtype=float32), 'flat_params_param': Array(0.01363992, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 33: log Z = -0.4630, l_oo = {'~': {'bias_scale_param': Array(0.01114203, dtype=float32), 'flat_params_param': Array(0.01363992, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 34: log Z = -0.4629, l_oo = {'~': {'bias_scale_param': Array(0.01062901, dtype=float32), 'flat_params_param': Array(0.00831449, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 34: log Z = -0.4629, l_oo = {'~': {'bias_scale_param': Array(0.01062901, dtype=float32), 'flat_params_param': Array(0.00831449, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 35: log Z = -0.4627, l_oo = {'~': {'bias_scale_param': Array(0.01019342, dtype=float32), 'flat_params_param': Array(0.016771, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 35: log Z = -0.4627, l_oo = {'~': {'bias_scale_param': Array(0.01019342, dtype=float32), 'flat_params_param': Array(0.016771, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 36: log Z = -0.4626, l_oo = {'~': {'bias_scale_param': Array(0.00999129, dtype=float32), 'flat_params_param': Array(0.02098086, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 36: log Z = -0.4626, l_oo = {'~': {'bias_scale_param': Array(0.00999129, dtype=float32), 'flat_params_param': Array(0.02098086, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 37: log Z = -0.4624, l_oo = {'~': {'bias_scale_param': Array(0.01004468, dtype=float32), 'flat_params_param': Array(0.01452026, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 37: log Z = -0.4624, l_oo = {'~': {'bias_scale_param': Array(0.01004468, dtype=float32), 'flat_params_param': Array(0.01452026, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 38: log Z = -0.4622, l_oo = {'~': {'bias_scale_param': Array(0.01024436, dtype=float32), 'flat_params_param': Array(0.00802776, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 38: log Z = -0.4622, l_oo = {'~': {'bias_scale_param': Array(0.01024436, dtype=float32), 'flat_params_param': Array(0.00802776, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 39: log Z = -0.4621, l_oo = {'~': {'bias_scale_param': Array(0.01041088, dtype=float32), 'flat_params_param': Array(0.01028205, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 39: log Z = -0.4621, l_oo = {'~': {'bias_scale_param': Array(0.01041088, dtype=float32), 'flat_params_param': Array(0.01028205, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 40: log Z = -0.4619, l_oo = {'~': {'bias_scale_param': Array(0.01039263, dtype=float32), 'flat_params_param': Array(0.014712, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 40: log Z = -0.4619, l_oo = {'~': {'bias_scale_param': Array(0.01039263, dtype=float32), 'flat_params_param': Array(0.014712, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 41: log Z = -0.4617, l_oo = {'~': {'bias_scale_param': Array(0.01015072, dtype=float32), 'flat_params_param': Array(0.0097152, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 41: log Z = -0.4617, l_oo = {'~': {'bias_scale_param': Array(0.01015072, dtype=float32), 'flat_params_param': Array(0.0097152, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 42: log Z = -0.4616, l_oo = {'~': {'bias_scale_param': Array(0.00977588, dtype=float32), 'flat_params_param': Array(0.00776856, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 42: log Z = -0.4616, l_oo = {'~': {'bias_scale_param': Array(0.00977588, dtype=float32), 'flat_params_param': Array(0.00776856, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 43: log Z = -0.4614, l_oo = {'~': {'bias_scale_param': Array(0.00942368, dtype=float32), 'flat_params_param': Array(0.01103474, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 43: log Z = -0.4614, l_oo = {'~': {'bias_scale_param': Array(0.00942368, dtype=float32), 'flat_params_param': Array(0.01103474, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 44: log Z = -0.4613, l_oo = {'~': {'bias_scale_param': Array(0.00922133, dtype=float32), 'flat_params_param': Array(0.01451756, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 44: log Z = -0.4613, l_oo = {'~': {'bias_scale_param': Array(0.00922133, dtype=float32), 'flat_params_param': Array(0.01451756, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 45: log Z = -0.4611, l_oo = {'~': {'bias_scale_param': Array(0.00920021, dtype=float32), 'flat_params_param': Array(0.00995263, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 45: log Z = -0.4611, l_oo = {'~': {'bias_scale_param': Array(0.00920021, dtype=float32), 'flat_params_param': Array(0.00995263, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 46: log Z = -0.4610, l_oo = {'~': {'bias_scale_param': Array(0.00928744, dtype=float32), 'flat_params_param': Array(0.00751892, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 46: log Z = -0.4610, l_oo = {'~': {'bias_scale_param': Array(0.00928744, dtype=float32), 'flat_params_param': Array(0.00751892, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 47: log Z = -0.4609, l_oo = {'~': {'bias_scale_param': Array(0.00935428, dtype=float32), 'flat_params_param': Array(0.00832261, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 47: log Z = -0.4609, l_oo = {'~': {'bias_scale_param': Array(0.00935428, dtype=float32), 'flat_params_param': Array(0.00832261, dtype=float32)}}:  27%|██▋       | 27/100 [00:00<00:01, 45.10it/s]
Step 47: log Z = -0.4609, l_oo = {'~': {'bias_scale_param': Array(0.00935428, dtype=float32), 'flat_params_param': Array(0.00832261, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 48: log Z = -0.4607, l_oo = {'~': {'bias_scale_param': Array(0.00929452, dtype=float32), 'flat_params_param': Array(0.01129548, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 48: log Z = -0.4607, l_oo = {'~': {'bias_scale_param': Array(0.00929452, dtype=float32), 'flat_params_param': Array(0.01129548, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 49: log Z = -0.4606, l_oo = {'~': {'bias_scale_param': Array(0.00908758, dtype=float32), 'flat_params_param': Array(0.00733504, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 49: log Z = -0.4606, l_oo = {'~': {'bias_scale_param': Array(0.00908758, dtype=float32), 'flat_params_param': Array(0.00733504, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 50: log Z = -0.4604, l_oo = {'~': {'bias_scale_param': Array(0.00880608, dtype=float32), 'flat_params_param': Array(0.00728814, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 50: log Z = -0.4604, l_oo = {'~': {'bias_scale_param': Array(0.00880608, dtype=float32), 'flat_params_param': Array(0.00728814, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 51: log Z = -0.4603, l_oo = {'~': {'bias_scale_param': Array(0.00856186, dtype=float32), 'flat_params_param': Array(0.00732028, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 51: log Z = -0.4603, l_oo = {'~': {'bias_scale_param': Array(0.00856186, dtype=float32), 'flat_params_param': Array(0.00732028, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 52: log Z = -0.4602, l_oo = {'~': {'bias_scale_param': Array(0.00843433, dtype=float32), 'flat_params_param': Array(0.00750329, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 52: log Z = -0.4602, l_oo = {'~': {'bias_scale_param': Array(0.00843433, dtype=float32), 'flat_params_param': Array(0.00750329, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 53: log Z = -0.4600, l_oo = {'~': {'bias_scale_param': Array(0.00842625, dtype=float32), 'flat_params_param': Array(0.00715561, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 53: log Z = -0.4600, l_oo = {'~': {'bias_scale_param': Array(0.00842625, dtype=float32), 'flat_params_param': Array(0.00715561, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 54: log Z = -0.4599, l_oo = {'~': {'bias_scale_param': Array(0.00846817, dtype=float32), 'flat_params_param': Array(0.00700348, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 54: log Z = -0.4599, l_oo = {'~': {'bias_scale_param': Array(0.00846817, dtype=float32), 'flat_params_param': Array(0.00700348, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 55: log Z = -0.4598, l_oo = {'~': {'bias_scale_param': Array(0.00846678, dtype=float32), 'flat_params_param': Array(0.00908117, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 55: log Z = -0.4598, l_oo = {'~': {'bias_scale_param': Array(0.00846678, dtype=float32), 'flat_params_param': Array(0.00908117, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 56: log Z = -0.4597, l_oo = {'~': {'bias_scale_param': Array(0.00836587, dtype=float32), 'flat_params_param': Array(0.0088526, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 56: log Z = -0.4597, l_oo = {'~': {'bias_scale_param': Array(0.00836587, dtype=float32), 'flat_params_param': Array(0.0088526, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 57: log Z = -0.4595, l_oo = {'~': {'bias_scale_param': Array(0.00818064, dtype=float32), 'flat_params_param': Array(0.00684582, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 57: log Z = -0.4595, l_oo = {'~': {'bias_scale_param': Array(0.00818064, dtype=float32), 'flat_params_param': Array(0.00684582, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 58: log Z = -0.4594, l_oo = {'~': {'bias_scale_param': Array(0.00798153, dtype=float32), 'flat_params_param': Array(0.0068695, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 58: log Z = -0.4594, l_oo = {'~': {'bias_scale_param': Array(0.00798153, dtype=float32), 'flat_params_param': Array(0.0068695, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 59: log Z = -0.4593, l_oo = {'~': {'bias_scale_param': Array(0.0078433, dtype=float32), 'flat_params_param': Array(0.00685484, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 59: log Z = -0.4593, l_oo = {'~': {'bias_scale_param': Array(0.0078433, dtype=float32), 'flat_params_param': Array(0.00685484, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 60: log Z = -0.4592, l_oo = {'~': {'bias_scale_param': Array(0.00779505, dtype=float32), 'flat_params_param': Array(0.00678139, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 60: log Z = -0.4592, l_oo = {'~': {'bias_scale_param': Array(0.00779505, dtype=float32), 'flat_params_param': Array(0.00678139, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 61: log Z = -0.4591, l_oo = {'~': {'bias_scale_param': Array(0.00780536, dtype=float32), 'flat_params_param': Array(0.00667022, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 61: log Z = -0.4591, l_oo = {'~': {'bias_scale_param': Array(0.00780536, dtype=float32), 'flat_params_param': Array(0.00667022, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 62: log Z = -0.4590, l_oo = {'~': {'bias_scale_param': Array(0.00780734, dtype=float32), 'flat_params_param': Array(0.00656663, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 62: log Z = -0.4590, l_oo = {'~': {'bias_scale_param': Array(0.00780734, dtype=float32), 'flat_params_param': Array(0.00656663, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 63: log Z = -0.4589, l_oo = {'~': {'bias_scale_param': Array(0.0077471, dtype=float32), 'flat_params_param': Array(0.00684452, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 63: log Z = -0.4589, l_oo = {'~': {'bias_scale_param': Array(0.0077471, dtype=float32), 'flat_params_param': Array(0.00684452, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 64: log Z = -0.4587, l_oo = {'~': {'bias_scale_param': Array(0.00761946, dtype=float32), 'flat_params_param': Array(0.006497, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 64: log Z = -0.4587, l_oo = {'~': {'bias_scale_param': Array(0.00761946, dtype=float32), 'flat_params_param': Array(0.006497, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 65: log Z = -0.4586, l_oo = {'~': {'bias_scale_param': Array(0.00746807, dtype=float32), 'flat_params_param': Array(0.00650447, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 65: log Z = -0.4586, l_oo = {'~': {'bias_scale_param': Array(0.00746807, dtype=float32), 'flat_params_param': Array(0.00650447, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 66: log Z = -0.4585, l_oo = {'~': {'bias_scale_param': Array(0.00735004, dtype=float32), 'flat_params_param': Array(0.00649015, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 66: log Z = -0.4585, l_oo = {'~': {'bias_scale_param': Array(0.00735004, dtype=float32), 'flat_params_param': Array(0.00649015, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 67: log Z = -0.4584, l_oo = {'~': {'bias_scale_param': Array(0.00729457, dtype=float32), 'flat_params_param': Array(0.00643406, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 67: log Z = -0.4584, l_oo = {'~': {'bias_scale_param': Array(0.00729457, dtype=float32), 'flat_params_param': Array(0.00643406, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 68: log Z = -0.4583, l_oo = {'~': {'bias_scale_param': Array(0.00728449, dtype=float32), 'flat_params_param': Array(0.0063478, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 68: log Z = -0.4583, l_oo = {'~': {'bias_scale_param': Array(0.00728449, dtype=float32), 'flat_params_param': Array(0.0063478, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 69: log Z = -0.4582, l_oo = {'~': {'bias_scale_param': Array(0.00727224, dtype=float32), 'flat_params_param': Array(0.00626371, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 69: log Z = -0.4582, l_oo = {'~': {'bias_scale_param': Array(0.00727224, dtype=float32), 'flat_params_param': Array(0.00626371, dtype=float32)}}:  48%|████▊     | 48/100 [00:00<00:00, 76.55it/s]
Step 70: log Z = -0.4581, l_oo = {'~': {'bias_scale_param': Array(0.00721649, dtype=float32), 'flat_params_param': Array(0.00621061, dtype=float32)}}:  48%|████▊     | 48/100 [00:01<00:00, 76.55it/s]
Step 70: log Z = -0.4581, l_oo = {'~': {'bias_scale_param': Array(0.00721649, dtype=float32), 'flat_params_param': Array(0.00621061, dtype=float32)}}:  48%|████▊     | 48/100 [00:01<00:00, 76.55it/s]
Step 70: log Z = -0.4581, l_oo = {'~': {'bias_scale_param': Array(0.00721649, dtype=float32), 'flat_params_param': Array(0.00621061, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 71: log Z = -0.4580, l_oo = {'~': {'bias_scale_param': Array(0.00711235, dtype=float32), 'flat_params_param': Array(0.0061918, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 71: log Z = -0.4580, l_oo = {'~': {'bias_scale_param': Array(0.00711235, dtype=float32), 'flat_params_param': Array(0.0061918, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 72: log Z = -0.4579, l_oo = {'~': {'bias_scale_param': Array(0.00699147, dtype=float32), 'flat_params_param': Array(0.0061855, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 72: log Z = -0.4579, l_oo = {'~': {'bias_scale_param': Array(0.00699147, dtype=float32), 'flat_params_param': Array(0.0061855, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 73: log Z = -0.4578, l_oo = {'~': {'bias_scale_param': Array(0.00689599, dtype=float32), 'flat_params_param': Array(0.00616169, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 73: log Z = -0.4578, l_oo = {'~': {'bias_scale_param': Array(0.00689599, dtype=float32), 'flat_params_param': Array(0.00616169, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 74: log Z = -0.4577, l_oo = {'~': {'bias_scale_param': Array(0.00684482, dtype=float32), 'flat_params_param': Array(0.00610796, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 74: log Z = -0.4577, l_oo = {'~': {'bias_scale_param': Array(0.00684482, dtype=float32), 'flat_params_param': Array(0.00610796, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 75: log Z = -0.4576, l_oo = {'~': {'bias_scale_param': Array(0.00682275, dtype=float32), 'flat_params_param': Array(0.00603442, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 75: log Z = -0.4576, l_oo = {'~': {'bias_scale_param': Array(0.00682275, dtype=float32), 'flat_params_param': Array(0.00603442, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 76: log Z = -0.4575, l_oo = {'~': {'bias_scale_param': Array(0.00679318, dtype=float32), 'flat_params_param': Array(0.00596667, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 76: log Z = -0.4575, l_oo = {'~': {'bias_scale_param': Array(0.00679318, dtype=float32), 'flat_params_param': Array(0.00596667, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 77: log Z = -0.4574, l_oo = {'~': {'bias_scale_param': Array(0.00672984, dtype=float32), 'flat_params_param': Array(0.00592326, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 77: log Z = -0.4574, l_oo = {'~': {'bias_scale_param': Array(0.00672984, dtype=float32), 'flat_params_param': Array(0.00592326, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 78: log Z = -0.4574, l_oo = {'~': {'bias_scale_param': Array(0.00663519, dtype=float32), 'flat_params_param': Array(0.00590249, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 78: log Z = -0.4574, l_oo = {'~': {'bias_scale_param': Array(0.00663519, dtype=float32), 'flat_params_param': Array(0.00590249, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 79: log Z = -0.4573, l_oo = {'~': {'bias_scale_param': Array(0.00653671, dtype=float32), 'flat_params_param': Array(0.00588476, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 79: log Z = -0.4573, l_oo = {'~': {'bias_scale_param': Array(0.00653671, dtype=float32), 'flat_params_param': Array(0.00588476, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 80: log Z = -0.4572, l_oo = {'~': {'bias_scale_param': Array(0.0064628, dtype=float32), 'flat_params_param': Array(0.00585077, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 80: log Z = -0.4572, l_oo = {'~': {'bias_scale_param': Array(0.0064628, dtype=float32), 'flat_params_param': Array(0.00585077, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 81: log Z = -0.4571, l_oo = {'~': {'bias_scale_param': Array(0.00642003, dtype=float32), 'flat_params_param': Array(0.00579608, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 81: log Z = -0.4571, l_oo = {'~': {'bias_scale_param': Array(0.00642003, dtype=float32), 'flat_params_param': Array(0.00579608, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 82: log Z = -0.4570, l_oo = {'~': {'bias_scale_param': Array(0.00639022, dtype=float32), 'flat_params_param': Array(0.00573299, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 82: log Z = -0.4570, l_oo = {'~': {'bias_scale_param': Array(0.00639022, dtype=float32), 'flat_params_param': Array(0.00573299, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 83: log Z = -0.4569, l_oo = {'~': {'bias_scale_param': Array(0.00634702, dtype=float32), 'flat_params_param': Array(0.0056803, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 83: log Z = -0.4569, l_oo = {'~': {'bias_scale_param': Array(0.00634702, dtype=float32), 'flat_params_param': Array(0.0056803, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 84: log Z = -0.4568, l_oo = {'~': {'bias_scale_param': Array(0.00627789, dtype=float32), 'flat_params_param': Array(0.00564685, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 84: log Z = -0.4568, l_oo = {'~': {'bias_scale_param': Array(0.00627789, dtype=float32), 'flat_params_param': Array(0.00564685, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 85: log Z = -0.4568, l_oo = {'~': {'bias_scale_param': Array(0.00619384, dtype=float32), 'flat_params_param': Array(0.00562517, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 85: log Z = -0.4568, l_oo = {'~': {'bias_scale_param': Array(0.00619384, dtype=float32), 'flat_params_param': Array(0.00562517, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 86: log Z = -0.4567, l_oo = {'~': {'bias_scale_param': Array(0.00611765, dtype=float32), 'flat_params_param': Array(0.00559856, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 86: log Z = -0.4567, l_oo = {'~': {'bias_scale_param': Array(0.00611765, dtype=float32), 'flat_params_param': Array(0.00559856, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 87: log Z = -0.4566, l_oo = {'~': {'bias_scale_param': Array(0.00606415, dtype=float32), 'flat_params_param': Array(0.00555742, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 87: log Z = -0.4566, l_oo = {'~': {'bias_scale_param': Array(0.00606415, dtype=float32), 'flat_params_param': Array(0.00555742, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 88: log Z = -0.4565, l_oo = {'~': {'bias_scale_param': Array(0.00602893, dtype=float32), 'flat_params_param': Array(0.00550484, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 88: log Z = -0.4565, l_oo = {'~': {'bias_scale_param': Array(0.00602893, dtype=float32), 'flat_params_param': Array(0.00550484, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 89: log Z = -0.4564, l_oo = {'~': {'bias_scale_param': Array(0.00599284, dtype=float32), 'flat_params_param': Array(0.0054539, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 89: log Z = -0.4564, l_oo = {'~': {'bias_scale_param': Array(0.00599284, dtype=float32), 'flat_params_param': Array(0.0054539, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 90: log Z = -0.4564, l_oo = {'~': {'bias_scale_param': Array(0.0059405, dtype=float32), 'flat_params_param': Array(0.00541559, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 90: log Z = -0.4564, l_oo = {'~': {'bias_scale_param': Array(0.0059405, dtype=float32), 'flat_params_param': Array(0.00541559, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 91: log Z = -0.4563, l_oo = {'~': {'bias_scale_param': Array(0.00587283, dtype=float32), 'flat_params_param': Array(0.00538884, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 91: log Z = -0.4563, l_oo = {'~': {'bias_scale_param': Array(0.00587283, dtype=float32), 'flat_params_param': Array(0.00538884, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 92: log Z = -0.4562, l_oo = {'~': {'bias_scale_param': Array(0.00580432, dtype=float32), 'flat_params_param': Array(0.00536431, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 92: log Z = -0.4562, l_oo = {'~': {'bias_scale_param': Array(0.00580432, dtype=float32), 'flat_params_param': Array(0.00536431, dtype=float32)}}:  71%|███████   | 71/100 [00:01<00:00, 109.46it/s]
Step 92: log Z = -0.4562, l_oo = {'~': {'bias_scale_param': Array(0.00580432, dtype=float32), 'flat_params_param': Array(0.00536431, dtype=float32)}}:  93%|█████████▎| 93/100 [00:01<00:00, 134.06it/s]
Step 93: log Z = -0.4561, l_oo = {'~': {'bias_scale_param': Array(0.00574974, dtype=float32), 'flat_params_param': Array(0.00533066, dtype=float32)}}:  93%|█████████▎| 93/100 [00:01<00:00, 134.06it/s]
Step 93: log Z = -0.4561, l_oo = {'~': {'bias_scale_param': Array(0.00574974, dtype=float32), 'flat_params_param': Array(0.00533066, dtype=float32)}}:  93%|█████████▎| 93/100 [00:01<00:00, 134.06it/s]
Step 94: log Z = -0.4561, l_oo = {'~': {'bias_scale_param': Array(0.00571117, dtype=float32), 'flat_params_param': Array(0.00528691, dtype=float32)}}:  93%|█████████▎| 93/100 [00:01<00:00, 134.06it/s]
Step 94: log Z = -0.4561, l_oo = {'~': {'bias_scale_param': Array(0.00571117, dtype=float32), 'flat_params_param': Array(0.00528691, dtype=float32)}}:  93%|█████████▎| 93/100 [00:01<00:00, 134.06it/s]
Step 95: log Z = -0.4560, l_oo = {'~': {'bias_scale_param': Array(0.00567714, dtype=float32), 'flat_params_param': Array(0.00524107, dtype=float32)}}:  93%|█████████▎| 93/100 [00:01<00:00, 134.06it/s]
Step 95: log Z = -0.4560, l_oo = {'~': {'bias_scale_param': Array(0.00567714, dtype=float32), 'flat_params_param': Array(0.00524107, dtype=float32)}}:  93%|█████████▎| 93/100 [00:01<00:00, 134.06it/s]
Step 96: log Z = -0.4559, l_oo = {'~': {'bias_scale_param': Array(0.005634, dtype=float32), 'flat_params_param': Array(0.00520258, dtype=float32)}}:  93%|█████████▎| 93/100 [00:01<00:00, 134.06it/s]
Step 96: log Z = -0.4559, l_oo = {'~': {'bias_scale_param': Array(0.005634, dtype=float32), 'flat_params_param': Array(0.00520258, dtype=float32)}}:  93%|█████████▎| 93/100 [00:01<00:00, 134.06it/s]
Step 97: log Z = -0.4558, l_oo = {'~': {'bias_scale_param': Array(0.00557844, dtype=float32), 'flat_params_param': Array(0.00517372, dtype=float32)}}:  93%|█████████▎| 93/100 [00:01<00:00, 134.06it/s]
Step 97: log Z = -0.4558, l_oo = {'~': {'bias_scale_param': Array(0.00557844, dtype=float32), 'flat_params_param': Array(0.00517372, dtype=float32)}}:  93%|█████████▎| 93/100 [00:01<00:00, 134.06it/s]
Step 98: log Z = -0.4558, l_oo = {'~': {'bias_scale_param': Array(0.00551883, dtype=float32), 'flat_params_param': Array(0.00514892, dtype=float32)}}:  93%|█████████▎| 93/100 [00:01<00:00, 134.06it/s]
Step 98: log Z = -0.4558, l_oo = {'~': {'bias_scale_param': Array(0.00551883, dtype=float32), 'flat_params_param': Array(0.00514892, dtype=float32)}}:  93%|█████████▎| 93/100 [00:01<00:00, 134.06it/s]
Step 99: log Z = -0.4557, l_oo = {'~': {'bias_scale_param': Array(0.00546731, dtype=float32), 'flat_params_param': Array(0.00511908, dtype=float32)}}:  93%|█████████▎| 93/100 [00:01<00:00, 134.06it/s]
Step 99: log Z = -0.4557, l_oo = {'~': {'bias_scale_param': Array(0.00546731, dtype=float32), 'flat_params_param': Array(0.00511908, dtype=float32)}}: 100%|██████████| 100/100 [00:01<00:00, 87.03it/s]
Step 6: log Z = -0.4512 +- 0.0171:  60%|██████    | 6/10 [00:14<00:07,  1.95s/it]
Training Progress:   0%|          | 0/100 [00:00<?, ?it/s]
Step 0: log Z = -0.4512, l_oo = {'~': {'bias_scale_param': Array(0.00329516, dtype=float32), 'flat_params_param': Array(0.02281182, dtype=float32)}}:   0%|          | 0/100 [00:00<?, ?it/s]
Step 0: log Z = -0.4512, l_oo = {'~': {'bias_scale_param': Array(0.00329516, dtype=float32), 'flat_params_param': Array(0.02281182, dtype=float32)}}:   0%|          | 0/100 [00:00<?, ?it/s]
Step 0: log Z = -0.4512, l_oo = {'~': {'bias_scale_param': Array(0.00329516, dtype=float32), 'flat_params_param': Array(0.02281182, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 1: log Z = -0.4582, l_oo = {'~': {'bias_scale_param': Array(0.00047503, dtype=float32), 'flat_params_param': Array(0.22835697, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 1: log Z = -0.4582, l_oo = {'~': {'bias_scale_param': Array(0.00047503, dtype=float32), 'flat_params_param': Array(0.22835697, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 2: log Z = -0.4515, l_oo = {'~': {'bias_scale_param': Array(0.00201886, dtype=float32), 'flat_params_param': Array(0.05647885, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 2: log Z = -0.4515, l_oo = {'~': {'bias_scale_param': Array(0.00201886, dtype=float32), 'flat_params_param': Array(0.05647885, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 3: log Z = -0.4528, l_oo = {'~': {'bias_scale_param': Array(0.00478138, dtype=float32), 'flat_params_param': Array(0.12832226, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 3: log Z = -0.4528, l_oo = {'~': {'bias_scale_param': Array(0.00478138, dtype=float32), 'flat_params_param': Array(0.12832226, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 4: log Z = -0.4550, l_oo = {'~': {'bias_scale_param': Array(0.00564039, dtype=float32), 'flat_params_param': Array(0.18859409, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 4: log Z = -0.4550, l_oo = {'~': {'bias_scale_param': Array(0.00564039, dtype=float32), 'flat_params_param': Array(0.18859409, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 5: log Z = -0.4529, l_oo = {'~': {'bias_scale_param': Array(0.00482026, dtype=float32), 'flat_params_param': Array(0.13899802, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 5: log Z = -0.4529, l_oo = {'~': {'bias_scale_param': Array(0.00482026, dtype=float32), 'flat_params_param': Array(0.13899802, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 6: log Z = -0.4509, l_oo = {'~': {'bias_scale_param': Array(0.00336446, dtype=float32), 'flat_params_param': Array(0.04496, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 6: log Z = -0.4509, l_oo = {'~': {'bias_scale_param': Array(0.00336446, dtype=float32), 'flat_params_param': Array(0.04496, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 7: log Z = -0.4510, l_oo = {'~': {'bias_scale_param': Array(0.00197756, dtype=float32), 'flat_params_param': Array(0.04811102, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 7: log Z = -0.4510, l_oo = {'~': {'bias_scale_param': Array(0.00197756, dtype=float32), 'flat_params_param': Array(0.04811102, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 8: log Z = -0.4523, l_oo = {'~': {'bias_scale_param': Array(0.00112207, dtype=float32), 'flat_params_param': Array(0.10716505, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 8: log Z = -0.4523, l_oo = {'~': {'bias_scale_param': Array(0.00112207, dtype=float32), 'flat_params_param': Array(0.10716505, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 9: log Z = -0.4526, l_oo = {'~': {'bias_scale_param': Array(0.00092764, dtype=float32), 'flat_params_param': Array(0.11994807, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 9: log Z = -0.4526, l_oo = {'~': {'bias_scale_param': Array(0.00092764, dtype=float32), 'flat_params_param': Array(0.11994807, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 10: log Z = -0.4517, l_oo = {'~': {'bias_scale_param': Array(0.00128379, dtype=float32), 'flat_params_param': Array(0.09259591, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 10: log Z = -0.4517, l_oo = {'~': {'bias_scale_param': Array(0.00128379, dtype=float32), 'flat_params_param': Array(0.09259591, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 11: log Z = -0.4506, l_oo = {'~': {'bias_scale_param': Array(0.0020041, dtype=float32), 'flat_params_param': Array(0.03850745, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 11: log Z = -0.4506, l_oo = {'~': {'bias_scale_param': Array(0.0020041, dtype=float32), 'flat_params_param': Array(0.03850745, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 12: log Z = -0.4503, l_oo = {'~': {'bias_scale_param': Array(0.00285704, dtype=float32), 'flat_params_param': Array(0.0252828, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 12: log Z = -0.4503, l_oo = {'~': {'bias_scale_param': Array(0.00285704, dtype=float32), 'flat_params_param': Array(0.0252828, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 13: log Z = -0.4509, l_oo = {'~': {'bias_scale_param': Array(0.00355877, dtype=float32), 'flat_params_param': Array(0.07838126, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 13: log Z = -0.4509, l_oo = {'~': {'bias_scale_param': Array(0.00355877, dtype=float32), 'flat_params_param': Array(0.07838126, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 14: log Z = -0.4513, l_oo = {'~': {'bias_scale_param': Array(0.00385937, dtype=float32), 'flat_params_param': Array(0.10290173, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 14: log Z = -0.4513, l_oo = {'~': {'bias_scale_param': Array(0.00385937, dtype=float32), 'flat_params_param': Array(0.10290173, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 15: log Z = -0.4510, l_oo = {'~': {'bias_scale_param': Array(0.00368808, dtype=float32), 'flat_params_param': Array(0.09328296, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 15: log Z = -0.4510, l_oo = {'~': {'bias_scale_param': Array(0.00368808, dtype=float32), 'flat_params_param': Array(0.09328296, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 16: log Z = -0.4503, l_oo = {'~': {'bias_scale_param': Array(0.00316391, dtype=float32), 'flat_params_param': Array(0.05724273, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 16: log Z = -0.4503, l_oo = {'~': {'bias_scale_param': Array(0.00316391, dtype=float32), 'flat_params_param': Array(0.05724273, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 17: log Z = -0.4500, l_oo = {'~': {'bias_scale_param': Array(0.00249759, dtype=float32), 'flat_params_param': Array(0.00944252, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 17: log Z = -0.4500, l_oo = {'~': {'bias_scale_param': Array(0.00249759, dtype=float32), 'flat_params_param': Array(0.00944252, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 18: log Z = -0.4501, l_oo = {'~': {'bias_scale_param': Array(0.00190344, dtype=float32), 'flat_params_param': Array(0.03428394, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 18: log Z = -0.4501, l_oo = {'~': {'bias_scale_param': Array(0.00190344, dtype=float32), 'flat_params_param': Array(0.03428394, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 19: log Z = -0.4504, l_oo = {'~': {'bias_scale_param': Array(0.0015371, dtype=float32), 'flat_params_param': Array(0.06145385, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 19: log Z = -0.4504, l_oo = {'~': {'bias_scale_param': Array(0.0015371, dtype=float32), 'flat_params_param': Array(0.06145385, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 20: log Z = -0.4504, l_oo = {'~': {'bias_scale_param': Array(0.00146174, dtype=float32), 'flat_params_param': Array(0.06619117, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 20: log Z = -0.4504, l_oo = {'~': {'bias_scale_param': Array(0.00146174, dtype=float32), 'flat_params_param': Array(0.06619117, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 21: log Z = -0.4501, l_oo = {'~': {'bias_scale_param': Array(0.00165517, dtype=float32), 'flat_params_param': Array(0.04955257, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 21: log Z = -0.4501, l_oo = {'~': {'bias_scale_param': Array(0.00165517, dtype=float32), 'flat_params_param': Array(0.04955257, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 22: log Z = -0.4497, l_oo = {'~': {'bias_scale_param': Array(0.00203188, dtype=float32), 'flat_params_param': Array(0.01903687, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 22: log Z = -0.4497, l_oo = {'~': {'bias_scale_param': Array(0.00203188, dtype=float32), 'flat_params_param': Array(0.01903687, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 23: log Z = -0.4496, l_oo = {'~': {'bias_scale_param': Array(0.00246226, dtype=float32), 'flat_params_param': Array(0.01799157, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 23: log Z = -0.4496, l_oo = {'~': {'bias_scale_param': Array(0.00246226, dtype=float32), 'flat_params_param': Array(0.01799157, dtype=float32)}}:   1%|          | 1/100 [00:00<01:00,  1.63it/s]
Step 23: log Z = -0.4496, l_oo = {'~': {'bias_scale_param': Array(0.00246226, dtype=float32), 'flat_params_param': Array(0.01799157, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 24: log Z = -0.4497, l_oo = {'~': {'bias_scale_param': Array(0.00279749, dtype=float32), 'flat_params_param': Array(0.04677823, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 24: log Z = -0.4497, l_oo = {'~': {'bias_scale_param': Array(0.00279749, dtype=float32), 'flat_params_param': Array(0.04677823, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 25: log Z = -0.4498, l_oo = {'~': {'bias_scale_param': Array(0.00292034, dtype=float32), 'flat_params_param': Array(0.05882759, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 25: log Z = -0.4498, l_oo = {'~': {'bias_scale_param': Array(0.00292034, dtype=float32), 'flat_params_param': Array(0.05882759, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 26: log Z = -0.4497, l_oo = {'~': {'bias_scale_param': Array(0.00279771, dtype=float32), 'flat_params_param': Array(0.05112609, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 26: log Z = -0.4497, l_oo = {'~': {'bias_scale_param': Array(0.00279771, dtype=float32), 'flat_params_param': Array(0.05112609, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 27: log Z = -0.4494, l_oo = {'~': {'bias_scale_param': Array(0.00248991, dtype=float32), 'flat_params_param': Array(0.02802431, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 27: log Z = -0.4494, l_oo = {'~': {'bias_scale_param': Array(0.00248991, dtype=float32), 'flat_params_param': Array(0.02802431, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 28: log Z = -0.4493, l_oo = {'~': {'bias_scale_param': Array(0.00211415, dtype=float32), 'flat_params_param': Array(0.00402021, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 28: log Z = -0.4493, l_oo = {'~': {'bias_scale_param': Array(0.00211415, dtype=float32), 'flat_params_param': Array(0.00402021, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 29: log Z = -0.4493, l_oo = {'~': {'bias_scale_param': Array(0.00179448, dtype=float32), 'flat_params_param': Array(0.02659385, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 29: log Z = -0.4493, l_oo = {'~': {'bias_scale_param': Array(0.00179448, dtype=float32), 'flat_params_param': Array(0.02659385, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 30: log Z = -0.4494, l_oo = {'~': {'bias_scale_param': Array(0.0016195, dtype=float32), 'flat_params_param': Array(0.04002494, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 30: log Z = -0.4494, l_oo = {'~': {'bias_scale_param': Array(0.0016195, dtype=float32), 'flat_params_param': Array(0.04002494, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 31: log Z = -0.4493, l_oo = {'~': {'bias_scale_param': Array(0.00162068, dtype=float32), 'flat_params_param': Array(0.03840987, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 31: log Z = -0.4493, l_oo = {'~': {'bias_scale_param': Array(0.00162068, dtype=float32), 'flat_params_param': Array(0.03840987, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 32: log Z = -0.4491, l_oo = {'~': {'bias_scale_param': Array(0.00177271, dtype=float32), 'flat_params_param': Array(0.02357039, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 32: log Z = -0.4491, l_oo = {'~': {'bias_scale_param': Array(0.00177271, dtype=float32), 'flat_params_param': Array(0.02357039, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 33: log Z = -0.4490, l_oo = {'~': {'bias_scale_param': Array(0.00200537, dtype=float32), 'flat_params_param': Array(0.00384698, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 33: log Z = -0.4490, l_oo = {'~': {'bias_scale_param': Array(0.00200537, dtype=float32), 'flat_params_param': Array(0.00384698, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 34: log Z = -0.4490, l_oo = {'~': {'bias_scale_param': Array(0.00222426, dtype=float32), 'flat_params_param': Array(0.01976335, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 34: log Z = -0.4490, l_oo = {'~': {'bias_scale_param': Array(0.00222426, dtype=float32), 'flat_params_param': Array(0.01976335, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 35: log Z = -0.4490, l_oo = {'~': {'bias_scale_param': Array(0.00234162, dtype=float32), 'flat_params_param': Array(0.03229064, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 35: log Z = -0.4490, l_oo = {'~': {'bias_scale_param': Array(0.00234162, dtype=float32), 'flat_params_param': Array(0.03229064, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 36: log Z = -0.4490, l_oo = {'~': {'bias_scale_param': Array(0.00231268, dtype=float32), 'flat_params_param': Array(0.0319814, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 36: log Z = -0.4490, l_oo = {'~': {'bias_scale_param': Array(0.00231268, dtype=float32), 'flat_params_param': Array(0.0319814, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 37: log Z = -0.4489, l_oo = {'~': {'bias_scale_param': Array(0.00215343, dtype=float32), 'flat_params_param': Array(0.01992878, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 37: log Z = -0.4489, l_oo = {'~': {'bias_scale_param': Array(0.00215343, dtype=float32), 'flat_params_param': Array(0.01992878, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 38: log Z = -0.4488, l_oo = {'~': {'bias_scale_param': Array(0.00192896, dtype=float32), 'flat_params_param': Array(0.00366015, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 38: log Z = -0.4488, l_oo = {'~': {'bias_scale_param': Array(0.00192896, dtype=float32), 'flat_params_param': Array(0.00366015, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 39: log Z = -0.4487, l_oo = {'~': {'bias_scale_param': Array(0.00172137, dtype=float32), 'flat_params_param': Array(0.01541603, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 39: log Z = -0.4487, l_oo = {'~': {'bias_scale_param': Array(0.00172137, dtype=float32), 'flat_params_param': Array(0.01541603, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 40: log Z = -0.4487, l_oo = {'~': {'bias_scale_param': Array(0.0015965, dtype=float32), 'flat_params_param': Array(0.02512758, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 40: log Z = -0.4487, l_oo = {'~': {'bias_scale_param': Array(0.0015965, dtype=float32), 'flat_params_param': Array(0.02512758, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 41: log Z = -0.4487, l_oo = {'~': {'bias_scale_param': Array(0.00158267, dtype=float32), 'flat_params_param': Array(0.02452877, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 41: log Z = -0.4487, l_oo = {'~': {'bias_scale_param': Array(0.00158267, dtype=float32), 'flat_params_param': Array(0.02452877, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 42: log Z = -0.4486, l_oo = {'~': {'bias_scale_param': Array(0.00166567, dtype=float32), 'flat_params_param': Array(0.01470131, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 42: log Z = -0.4486, l_oo = {'~': {'bias_scale_param': Array(0.00166567, dtype=float32), 'flat_params_param': Array(0.01470131, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 43: log Z = -0.4485, l_oo = {'~': {'bias_scale_param': Array(0.00179735, dtype=float32), 'flat_params_param': Array(0.00351856, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 43: log Z = -0.4485, l_oo = {'~': {'bias_scale_param': Array(0.00179735, dtype=float32), 'flat_params_param': Array(0.00351856, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 44: log Z = -0.4485, l_oo = {'~': {'bias_scale_param': Array(0.00191376, dtype=float32), 'flat_params_param': Array(0.0132911, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 44: log Z = -0.4485, l_oo = {'~': {'bias_scale_param': Array(0.00191376, dtype=float32), 'flat_params_param': Array(0.0132911, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 45: log Z = -0.4485, l_oo = {'~': {'bias_scale_param': Array(0.00196039, dtype=float32), 'flat_params_param': Array(0.02008083, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 45: log Z = -0.4485, l_oo = {'~': {'bias_scale_param': Array(0.00196039, dtype=float32), 'flat_params_param': Array(0.02008083, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 46: log Z = -0.4484, l_oo = {'~': {'bias_scale_param': Array(0.00191509, dtype=float32), 'flat_params_param': Array(0.01798797, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 46: log Z = -0.4484, l_oo = {'~': {'bias_scale_param': Array(0.00191509, dtype=float32), 'flat_params_param': Array(0.01798797, dtype=float32)}}:  24%|██▍       | 24/100 [00:00<00:01, 44.61it/s]
Step 46: log Z = -0.4484, l_oo = {'~': {'bias_scale_param': Array(0.00191509, dtype=float32), 'flat_params_param': Array(0.01798797, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 47: log Z = -0.4484, l_oo = {'~': {'bias_scale_param': Array(0.001797, dtype=float32), 'flat_params_param': Array(0.0086513, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 47: log Z = -0.4484, l_oo = {'~': {'bias_scale_param': Array(0.001797, dtype=float32), 'flat_params_param': Array(0.0086513, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 48: log Z = -0.4483, l_oo = {'~': {'bias_scale_param': Array(0.00165369, dtype=float32), 'flat_params_param': Array(0.00341245, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 48: log Z = -0.4483, l_oo = {'~': {'bias_scale_param': Array(0.00165369, dtype=float32), 'flat_params_param': Array(0.00341245, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 49: log Z = -0.4483, l_oo = {'~': {'bias_scale_param': Array(0.00153816, dtype=float32), 'flat_params_param': Array(0.01293992, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 49: log Z = -0.4483, l_oo = {'~': {'bias_scale_param': Array(0.00153816, dtype=float32), 'flat_params_param': Array(0.01293992, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 50: log Z = -0.4483, l_oo = {'~': {'bias_scale_param': Array(0.00148605, dtype=float32), 'flat_params_param': Array(0.01621115, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 50: log Z = -0.4483, l_oo = {'~': {'bias_scale_param': Array(0.00148605, dtype=float32), 'flat_params_param': Array(0.01621115, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 51: log Z = -0.4482, l_oo = {'~': {'bias_scale_param': Array(0.00150399, dtype=float32), 'flat_params_param': Array(0.01237001, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 51: log Z = -0.4482, l_oo = {'~': {'bias_scale_param': Array(0.00150399, dtype=float32), 'flat_params_param': Array(0.01237001, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 52: log Z = -0.4481, l_oo = {'~': {'bias_scale_param': Array(0.00156958, dtype=float32), 'flat_params_param': Array(0.00357253, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 52: log Z = -0.4481, l_oo = {'~': {'bias_scale_param': Array(0.00156958, dtype=float32), 'flat_params_param': Array(0.00357253, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 53: log Z = -0.4481, l_oo = {'~': {'bias_scale_param': Array(0.00164231, dtype=float32), 'flat_params_param': Array(0.00606205, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 53: log Z = -0.4481, l_oo = {'~': {'bias_scale_param': Array(0.00164231, dtype=float32), 'flat_params_param': Array(0.00606205, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 54: log Z = -0.4481, l_oo = {'~': {'bias_scale_param': Array(0.00168157, dtype=float32), 'flat_params_param': Array(0.0122922, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 54: log Z = -0.4481, l_oo = {'~': {'bias_scale_param': Array(0.00168157, dtype=float32), 'flat_params_param': Array(0.0122922, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 55: log Z = -0.4480, l_oo = {'~': {'bias_scale_param': Array(0.00166473, dtype=float32), 'flat_params_param': Array(0.01265251, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 55: log Z = -0.4480, l_oo = {'~': {'bias_scale_param': Array(0.00166473, dtype=float32), 'flat_params_param': Array(0.01265251, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 56: log Z = -0.4480, l_oo = {'~': {'bias_scale_param': Array(0.00159626, dtype=float32), 'flat_params_param': Array(0.00747458, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 56: log Z = -0.4480, l_oo = {'~': {'bias_scale_param': Array(0.00159626, dtype=float32), 'flat_params_param': Array(0.00747458, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 57: log Z = -0.4480, l_oo = {'~': {'bias_scale_param': Array(0.00150414, dtype=float32), 'flat_params_param': Array(0.00313549, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 57: log Z = -0.4480, l_oo = {'~': {'bias_scale_param': Array(0.00150414, dtype=float32), 'flat_params_param': Array(0.00313549, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 58: log Z = -0.4479, l_oo = {'~': {'bias_scale_param': Array(0.00142419, dtype=float32), 'flat_params_param': Array(0.00711498, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 58: log Z = -0.4479, l_oo = {'~': {'bias_scale_param': Array(0.00142419, dtype=float32), 'flat_params_param': Array(0.00711498, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 59: log Z = -0.4479, l_oo = {'~': {'bias_scale_param': Array(0.00138289, dtype=float32), 'flat_params_param': Array(0.00976647, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 59: log Z = -0.4479, l_oo = {'~': {'bias_scale_param': Array(0.00138289, dtype=float32), 'flat_params_param': Array(0.00976647, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 60: log Z = -0.4479, l_oo = {'~': {'bias_scale_param': Array(0.00138714, dtype=float32), 'flat_params_param': Array(0.00747507, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 60: log Z = -0.4479, l_oo = {'~': {'bias_scale_param': Array(0.00138714, dtype=float32), 'flat_params_param': Array(0.00747507, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 61: log Z = -0.4478, l_oo = {'~': {'bias_scale_param': Array(0.00142305, dtype=float32), 'flat_params_param': Array(0.00304481, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 61: log Z = -0.4478, l_oo = {'~': {'bias_scale_param': Array(0.00142305, dtype=float32), 'flat_params_param': Array(0.00304481, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 62: log Z = -0.4478, l_oo = {'~': {'bias_scale_param': Array(0.00146353, dtype=float32), 'flat_params_param': Array(0.00468033, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 62: log Z = -0.4478, l_oo = {'~': {'bias_scale_param': Array(0.00146353, dtype=float32), 'flat_params_param': Array(0.00468033, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 63: log Z = -0.4477, l_oo = {'~': {'bias_scale_param': Array(0.00148151, dtype=float32), 'flat_params_param': Array(0.00854224, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 63: log Z = -0.4477, l_oo = {'~': {'bias_scale_param': Array(0.00148151, dtype=float32), 'flat_params_param': Array(0.00854224, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 64: log Z = -0.4477, l_oo = {'~': {'bias_scale_param': Array(0.00146317, dtype=float32), 'flat_params_param': Array(0.00829166, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 64: log Z = -0.4477, l_oo = {'~': {'bias_scale_param': Array(0.00146317, dtype=float32), 'flat_params_param': Array(0.00829166, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 65: log Z = -0.4477, l_oo = {'~': {'bias_scale_param': Array(0.0014134, dtype=float32), 'flat_params_param': Array(0.00439684, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 65: log Z = -0.4477, l_oo = {'~': {'bias_scale_param': Array(0.0014134, dtype=float32), 'flat_params_param': Array(0.00439684, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 66: log Z = -0.4476, l_oo = {'~': {'bias_scale_param': Array(0.00135212, dtype=float32), 'flat_params_param': Array(0.00292372, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 66: log Z = -0.4476, l_oo = {'~': {'bias_scale_param': Array(0.00135212, dtype=float32), 'flat_params_param': Array(0.00292372, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 67: log Z = -0.4476, l_oo = {'~': {'bias_scale_param': Array(0.00130294, dtype=float32), 'flat_params_param': Array(0.00498256, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 67: log Z = -0.4476, l_oo = {'~': {'bias_scale_param': Array(0.00130294, dtype=float32), 'flat_params_param': Array(0.00498256, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 68: log Z = -0.4476, l_oo = {'~': {'bias_scale_param': Array(0.00128139, dtype=float32), 'flat_params_param': Array(0.00589171, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 68: log Z = -0.4476, l_oo = {'~': {'bias_scale_param': Array(0.00128139, dtype=float32), 'flat_params_param': Array(0.00589171, dtype=float32)}}:  47%|████▋     | 47/100 [00:00<00:00, 83.38it/s]
Step 68: log Z = -0.4476, l_oo = {'~': {'bias_scale_param': Array(0.00128139, dtype=float32), 'flat_params_param': Array(0.00589171, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 69: log Z = -0.4475, l_oo = {'~': {'bias_scale_param': Array(0.00128824, dtype=float32), 'flat_params_param': Array(0.00356239, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 69: log Z = -0.4475, l_oo = {'~': {'bias_scale_param': Array(0.00128824, dtype=float32), 'flat_params_param': Array(0.00356239, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 70: log Z = -0.4475, l_oo = {'~': {'bias_scale_param': Array(0.00131046, dtype=float32), 'flat_params_param': Array(0.00282606, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 70: log Z = -0.4475, l_oo = {'~': {'bias_scale_param': Array(0.00131046, dtype=float32), 'flat_params_param': Array(0.00282606, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 71: log Z = -0.4475, l_oo = {'~': {'bias_scale_param': Array(0.00132883, dtype=float32), 'flat_params_param': Array(0.00451479, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 71: log Z = -0.4475, l_oo = {'~': {'bias_scale_param': Array(0.00132883, dtype=float32), 'flat_params_param': Array(0.00451479, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 72: log Z = -0.4475, l_oo = {'~': {'bias_scale_param': Array(0.00132764, dtype=float32), 'flat_params_param': Array(0.00596289, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 72: log Z = -0.4475, l_oo = {'~': {'bias_scale_param': Array(0.00132764, dtype=float32), 'flat_params_param': Array(0.00596289, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 73: log Z = -0.4474, l_oo = {'~': {'bias_scale_param': Array(0.00130267, dtype=float32), 'flat_params_param': Array(0.00451513, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 73: log Z = -0.4474, l_oo = {'~': {'bias_scale_param': Array(0.00130267, dtype=float32), 'flat_params_param': Array(0.00451513, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 74: log Z = -0.4474, l_oo = {'~': {'bias_scale_param': Array(0.00126254, dtype=float32), 'flat_params_param': Array(0.00273965, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 74: log Z = -0.4474, l_oo = {'~': {'bias_scale_param': Array(0.00126254, dtype=float32), 'flat_params_param': Array(0.00273965, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 75: log Z = -0.4474, l_oo = {'~': {'bias_scale_param': Array(0.00122287, dtype=float32), 'flat_params_param': Array(0.00273725, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 75: log Z = -0.4474, l_oo = {'~': {'bias_scale_param': Array(0.00122287, dtype=float32), 'flat_params_param': Array(0.00273725, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 76: log Z = -0.4473, l_oo = {'~': {'bias_scale_param': Array(0.0011976, dtype=float32), 'flat_params_param': Array(0.00387547, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 76: log Z = -0.4473, l_oo = {'~': {'bias_scale_param': Array(0.0011976, dtype=float32), 'flat_params_param': Array(0.00387547, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 77: log Z = -0.4473, l_oo = {'~': {'bias_scale_param': Array(0.00119205, dtype=float32), 'flat_params_param': Array(0.00311916, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 77: log Z = -0.4473, l_oo = {'~': {'bias_scale_param': Array(0.00119205, dtype=float32), 'flat_params_param': Array(0.00311916, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 78: log Z = -0.4473, l_oo = {'~': {'bias_scale_param': Array(0.00120091, dtype=float32), 'flat_params_param': Array(0.00266846, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 78: log Z = -0.4473, l_oo = {'~': {'bias_scale_param': Array(0.00120091, dtype=float32), 'flat_params_param': Array(0.00266846, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 79: log Z = -0.4473, l_oo = {'~': {'bias_scale_param': Array(0.00121172, dtype=float32), 'flat_params_param': Array(0.00263425, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 79: log Z = -0.4473, l_oo = {'~': {'bias_scale_param': Array(0.00121172, dtype=float32), 'flat_params_param': Array(0.00263425, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 80: log Z = -0.4472, l_oo = {'~': {'bias_scale_param': Array(0.00121239, dtype=float32), 'flat_params_param': Array(0.00379457, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 80: log Z = -0.4472, l_oo = {'~': {'bias_scale_param': Array(0.00121239, dtype=float32), 'flat_params_param': Array(0.00379457, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 81: log Z = -0.4472, l_oo = {'~': {'bias_scale_param': Array(0.00119727, dtype=float32), 'flat_params_param': Array(0.00331065, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 81: log Z = -0.4472, l_oo = {'~': {'bias_scale_param': Array(0.00119727, dtype=float32), 'flat_params_param': Array(0.00331065, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 82: log Z = -0.4472, l_oo = {'~': {'bias_scale_param': Array(0.0011701, dtype=float32), 'flat_params_param': Array(0.00258237, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 82: log Z = -0.4472, l_oo = {'~': {'bias_scale_param': Array(0.0011701, dtype=float32), 'flat_params_param': Array(0.00258237, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 83: log Z = -0.4471, l_oo = {'~': {'bias_scale_param': Array(0.00114068, dtype=float32), 'flat_params_param': Array(0.002576, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 83: log Z = -0.4471, l_oo = {'~': {'bias_scale_param': Array(0.00114068, dtype=float32), 'flat_params_param': Array(0.002576, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 84: log Z = -0.4471, l_oo = {'~': {'bias_scale_param': Array(0.00111937, dtype=float32), 'flat_params_param': Array(0.00256436, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 84: log Z = -0.4471, l_oo = {'~': {'bias_scale_param': Array(0.00111937, dtype=float32), 'flat_params_param': Array(0.00256436, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 85: log Z = -0.4471, l_oo = {'~': {'bias_scale_param': Array(0.00111107, dtype=float32), 'flat_params_param': Array(0.00254464, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 85: log Z = -0.4471, l_oo = {'~': {'bias_scale_param': Array(0.00111107, dtype=float32), 'flat_params_param': Array(0.00254464, dtype=float32)}}:  69%|██████▉   | 69/100 [00:00<00:00, 114.72it/s]
Step 86: log Z = -0.4471, l_oo = {'~': {'bias_scale_param': Array(0.0011132, dtype=float32), 'flat_params_param': Array(0.00251838, dtype=float32)}}:  69%|██████▉   | 69/100 [00:01<00:00, 114.72it/s]
Step 86: log Z = -0.4471, l_oo = {'~': {'bias_scale_param': Array(0.0011132, dtype=float32), 'flat_params_param': Array(0.00251838, dtype=float32)}}:  69%|██████▉   | 69/100 [00:01<00:00, 114.72it/s]
Step 87: log Z = -0.4470, l_oo = {'~': {'bias_scale_param': Array(0.00111754, dtype=float32), 'flat_params_param': Array(0.00249084, dtype=float32)}}:  69%|██████▉   | 69/100 [00:01<00:00, 114.72it/s]
Step 87: log Z = -0.4470, l_oo = {'~': {'bias_scale_param': Array(0.00111754, dtype=float32), 'flat_params_param': Array(0.00249084, dtype=float32)}}:  69%|██████▉   | 69/100 [00:01<00:00, 114.72it/s]
Step 88: log Z = -0.4470, l_oo = {'~': {'bias_scale_param': Array(0.00111572, dtype=float32), 'flat_params_param': Array(0.00247267, dtype=float32)}}:  69%|██████▉   | 69/100 [00:01<00:00, 114.72it/s]
Step 88: log Z = -0.4470, l_oo = {'~': {'bias_scale_param': Array(0.00111572, dtype=float32), 'flat_params_param': Array(0.00247267, dtype=float32)}}:  69%|██████▉   | 69/100 [00:01<00:00, 114.72it/s]
Step 89: log Z = -0.4470, l_oo = {'~': {'bias_scale_param': Array(0.00110348, dtype=float32), 'flat_params_param': Array(0.00245166, dtype=float32)}}:  69%|██████▉   | 69/100 [00:01<00:00, 114.72it/s]
Step 89: log Z = -0.4470, l_oo = {'~': {'bias_scale_param': Array(0.00110348, dtype=float32), 'flat_params_param': Array(0.00245166, dtype=float32)}}:  69%|██████▉   | 69/100 [00:01<00:00, 114.72it/s]
Step 90: log Z = -0.4470, l_oo = {'~': {'bias_scale_param': Array(0.00108314, dtype=float32), 'flat_params_param': Array(0.00244125, dtype=float32)}}:  69%|██████▉   | 69/100 [00:01<00:00, 114.72it/s]
Step 90: log Z = -0.4470, l_oo = {'~': {'bias_scale_param': Array(0.00108314, dtype=float32), 'flat_params_param': Array(0.00244125, dtype=float32)}}:  69%|██████▉   | 69/100 [00:01<00:00, 114.72it/s]
Step 90: log Z = -0.4470, l_oo = {'~': {'bias_scale_param': Array(0.00108314, dtype=float32), 'flat_params_param': Array(0.00244125, dtype=float32)}}:  91%|█████████ | 91/100 [00:01<00:00, 140.18it/s]
Step 91: log Z = -0.4469, l_oo = {'~': {'bias_scale_param': Array(0.0010614, dtype=float32), 'flat_params_param': Array(0.00243188, dtype=float32)}}:  91%|█████████ | 91/100 [00:01<00:00, 140.18it/s]
Step 91: log Z = -0.4469, l_oo = {'~': {'bias_scale_param': Array(0.0010614, dtype=float32), 'flat_params_param': Array(0.00243188, dtype=float32)}}:  91%|█████████ | 91/100 [00:01<00:00, 140.18it/s]
Step 92: log Z = -0.4469, l_oo = {'~': {'bias_scale_param': Array(0.00104531, dtype=float32), 'flat_params_param': Array(0.00241916, dtype=float32)}}:  91%|█████████ | 91/100 [00:01<00:00, 140.18it/s]
Step 92: log Z = -0.4469, l_oo = {'~': {'bias_scale_param': Array(0.00104531, dtype=float32), 'flat_params_param': Array(0.00241916, dtype=float32)}}:  91%|█████████ | 91/100 [00:01<00:00, 140.18it/s]
Step 93: log Z = -0.4469, l_oo = {'~': {'bias_scale_param': Array(0.00103793, dtype=float32), 'flat_params_param': Array(0.00240105, dtype=float32)}}:  91%|█████████ | 91/100 [00:01<00:00, 140.18it/s]
Step 93: log Z = -0.4469, l_oo = {'~': {'bias_scale_param': Array(0.00103793, dtype=float32), 'flat_params_param': Array(0.00240105, dtype=float32)}}:  91%|█████████ | 91/100 [00:01<00:00, 140.18it/s]
Step 94: log Z = -0.4469, l_oo = {'~': {'bias_scale_param': Array(0.00103715, dtype=float32), 'flat_params_param': Array(0.00237873, dtype=float32)}}:  91%|█████████ | 91/100 [00:01<00:00, 140.18it/s]
Step 94: log Z = -0.4469, l_oo = {'~': {'bias_scale_param': Array(0.00103715, dtype=float32), 'flat_params_param': Array(0.00237873, dtype=float32)}}:  91%|█████████ | 91/100 [00:01<00:00, 140.18it/s]
Step 95: log Z = -0.4469, l_oo = {'~': {'bias_scale_param': Array(0.0010372, dtype=float32), 'flat_params_param': Array(0.00235629, dtype=float32)}}:  91%|█████████ | 91/100 [00:01<00:00, 140.18it/s]
Step 95: log Z = -0.4469, l_oo = {'~': {'bias_scale_param': Array(0.0010372, dtype=float32), 'flat_params_param': Array(0.00235629, dtype=float32)}}:  91%|█████████ | 91/100 [00:01<00:00, 140.18it/s]
Step 96: log Z = -0.4468, l_oo = {'~': {'bias_scale_param': Array(0.00103252, dtype=float32), 'flat_params_param': Array(0.00233727, dtype=float32)}}:  91%|█████████ | 91/100 [00:01<00:00, 140.18it/s]
Step 96: log Z = -0.4468, l_oo = {'~': {'bias_scale_param': Array(0.00103252, dtype=float32), 'flat_params_param': Array(0.00233727, dtype=float32)}}:  91%|█████████ | 91/100 [00:01<00:00, 140.18it/s]
Step 97: log Z = -0.4468, l_oo = {'~': {'bias_scale_param': Array(0.00102097, dtype=float32), 'flat_params_param': Array(0.00232278, dtype=float32)}}:  91%|█████████ | 91/100 [00:01<00:00, 140.18it/s]
Step 97: log Z = -0.4468, l_oo = {'~': {'bias_scale_param': Array(0.00102097, dtype=float32), 'flat_params_param': Array(0.00232278, dtype=float32)}}:  91%|█████████ | 91/100 [00:01<00:00, 140.18it/s]
Step 98: log Z = -0.4468, l_oo = {'~': {'bias_scale_param': Array(0.00100473, dtype=float32), 'flat_params_param': Array(0.00231167, dtype=float32)}}:  91%|█████████ | 91/100 [00:01<00:00, 140.18it/s]
Step 98: log Z = -0.4468, l_oo = {'~': {'bias_scale_param': Array(0.00100473, dtype=float32), 'flat_params_param': Array(0.00231167, dtype=float32)}}:  91%|█████████ | 91/100 [00:01<00:00, 140.18it/s]
Step 99: log Z = -0.4468, l_oo = {'~': {'bias_scale_param': Array(0.00098859, dtype=float32), 'flat_params_param': Array(0.00230066, dtype=float32)}}:  91%|█████████ | 91/100 [00:01<00:00, 140.18it/s]
Step 99: log Z = -0.4468, l_oo = {'~': {'bias_scale_param': Array(0.00098859, dtype=float32), 'flat_params_param': Array(0.00230066, dtype=float32)}}: 100%|██████████| 100/100 [00:01<00:00, 93.80it/s]
Step 7: log Z = -0.4557 +- 0.0173:  70%|███████   | 7/10 [00:15<00:06,  2.24s/it]
Convergence achieved at step 7.

[9]:
from jaxns import summary, plot_diagnostics, plot_cornerplot

summary(results)
plot_diagnostics(results)
plot_cornerplot(results)
--------
Termination Conditions:
Small remaining evidence
--------
likelihood evals: 11628
samples: 360
phantom samples: 0
likelihood evals / sample: 32.3
phantom fraction (%): 0.0%
--------
logZ=-0.456 +- 0.017
H=-0.02
ESS=32
--------
bias: mean +- std.dev. | 10%ile / 50%ile / 90%ile | MAP est. | max(L) est.
bias: -0.0 +- 0.15 | -0.2 / 0.0 / 0.19 | 0.0 | 0.29
--------
../_images/examples_Bayesian_neural_network_9_1.png
../_images/examples_Bayesian_neural_network_9_2.png
[10]:
U_samples = resample(random.PRNGKey(45),
                     results.U_samples,
                     results.log_dp_mean,
                     S=int(results.ESS),
                     replace=True)


def predict(U_sample):
    # Apply the opt_params to the model before computing the logits
    logits = model(opt_params).prepare_input(U_sample)[0]
    return jax.nn.sigmoid(logits)


predictions = vmap(predict)(U_samples)
print(predictions)
for i in range(len(y)):
    preds = jnp.percentile(predictions[:, i], jnp.asarray([5, 50, 95]))
    hit = jnp.mean(predictions[:, i]) > 0.5
    print(f"{i}: {x[i]} -> {y[i]} | pred: {preds} {'✓' if hit == y[i] else '✗'}")

accuracy = jnp.mean((jnp.mean(predictions, axis=0) > 0.5) == y)
print(f"Accuracy: {accuracy * 100:.1f}%")
[[[0.81260693]
  [0.7820038 ]
  [0.7820038 ]
  ...
  [0.13765268]
  [0.13765268]
  [0.0493108 ]]

 [[0.81416523]
  [0.7846781 ]
  [0.7846781 ]
  ...
  [0.14661536]
  [0.14661536]
  [0.05315876]]

 [[0.80607915]
  [0.7707676 ]
  [0.7707676 ]
  ...
  [0.10695916]
  [0.10695916]
  [0.0368626 ]]

 ...

 [[0.8094097 ]
  [0.77650696]
  [0.77650696]
  ...
  [0.1213586 ]
  [0.1213586 ]
  [0.04256485]]

 [[0.8029672 ]
  [0.76539457]
  [0.76539457]
  ...
  [0.09548252]
  [0.09548252]
  [0.03248575]]

 [[0.8043955 ]
  [0.7678619 ]
  [0.7678619 ]
  ...
  [0.10053474]
  [0.10053474]
  [0.03439451]]]
0: [1. 1. 1. 1. 1. 1. 1.] -> [ True] | pred: [0.8030733 0.8095837 0.8212578] ✓
1: [1. 1. 1. 1. 1. 1. 0.] -> [ True] | pred: [0.76557785 0.7768066  0.7967997 ] ✓
2: [1. 1. 1. 1. 1. 0. 1.] -> [ True] | pred: [0.76557785 0.7768066  0.7967997 ] ✓
3: [1. 1. 1. 1. 1. 0. 0.] -> [ True] | pred: [0.7004969  0.7199961  0.75471765] ✓
4: [1. 1. 1. 1. 0. 1. 1.] -> [ True] | pred: [0.76557785 0.7768066  0.7967997 ] ✓
5: [1. 1. 1. 1. 0. 1. 0.] -> [ True] | pred: [0.7004969  0.7199961  0.75471765] ✓
6: [1. 1. 1. 1. 0. 0. 1.] -> [ True] | pred: [0.7004969  0.7199961  0.75471765] ✓
7: [1. 1. 1. 1. 0. 0. 0.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
8: [1. 1. 1. 0. 1. 1. 1.] -> [ True] | pred: [0.76557785 0.7768066  0.7967997 ] ✓
9: [1. 1. 1. 0. 1. 1. 0.] -> [ True] | pred: [0.7004969  0.7199961  0.75471765] ✓
10: [1. 1. 1. 0. 1. 0. 1.] -> [ True] | pred: [0.7004969  0.7199961  0.75471765] ✓
11: [1. 1. 1. 0. 1. 0. 0.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
12: [1. 1. 1. 0. 0. 1. 1.] -> [ True] | pred: [0.7004969  0.7199961  0.75471765] ✓
13: [1. 1. 1. 0. 0. 1. 0.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
14: [1. 1. 1. 0. 0. 0. 1.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
15: [1. 1. 1. 0. 0. 0. 0.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
16: [1. 1. 0. 1. 1. 1. 1.] -> [ True] | pred: [0.76557785 0.7768066  0.7967997 ] ✓
17: [1. 1. 0. 1. 1. 1. 0.] -> [ True] | pred: [0.7004969  0.7199961  0.75471765] ✓
18: [1. 1. 0. 1. 1. 0. 1.] -> [ True] | pred: [0.7004969  0.7199961  0.75471765] ✓
19: [1. 1. 0. 1. 1. 0. 0.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
20: [1. 1. 0. 1. 0. 1. 1.] -> [ True] | pred: [0.7004969  0.7199961  0.75471765] ✓
21: [1. 1. 0. 1. 0. 1. 0.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
22: [1. 1. 0. 1. 0. 0. 1.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
23: [1. 1. 0. 1. 0. 0. 0.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
24: [1. 1. 0. 0. 1. 1. 1.] -> [ True] | pred: [0.7004969  0.7199961  0.75471765] ✓
25: [1. 1. 0. 0. 1. 1. 0.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
26: [1. 1. 0. 0. 1. 0. 1.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
27: [1. 1. 0. 0. 1. 0. 0.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
28: [1. 1. 0. 0. 0. 1. 1.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
29: [1. 1. 0. 0. 0. 1. 0.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
30: [1. 1. 0. 0. 0. 0. 1.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
31: [1. 1. 0. 0. 0. 0. 0.] -> [False] | pred: [0.23323378 0.27784103 0.38287282] ✓
32: [1. 0. 1. 1. 1. 1. 1.] -> [ True] | pred: [0.76557785 0.7768066  0.7967997 ] ✓
33: [1. 0. 1. 1. 1. 1. 0.] -> [ True] | pred: [0.7004969  0.7199961  0.75471765] ✓
34: [1. 0. 1. 1. 1. 0. 1.] -> [ True] | pred: [0.7004969  0.7199961  0.75471765] ✓
35: [1. 0. 1. 1. 1. 0. 0.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
36: [1. 0. 1. 1. 0. 1. 1.] -> [ True] | pred: [0.7004969  0.7199961  0.75471765] ✓
37: [1. 0. 1. 1. 0. 1. 0.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
38: [1. 0. 1. 1. 0. 0. 1.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
39: [1. 0. 1. 1. 0. 0. 0.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
40: [1. 0. 1. 0. 1. 1. 1.] -> [ True] | pred: [0.7004969  0.7199961  0.75471765] ✓
41: [1. 0. 1. 0. 1. 1. 0.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
42: [1. 0. 1. 0. 1. 0. 1.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
43: [1. 0. 1. 0. 1. 0. 0.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
44: [1. 0. 1. 0. 0. 1. 1.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
45: [1. 0. 1. 0. 0. 1. 0.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
46: [1. 0. 1. 0. 0. 0. 1.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
47: [1. 0. 1. 0. 0. 0. 0.] -> [False] | pred: [0.23323378 0.27784103 0.38287282] ✓
48: [1. 0. 0. 1. 1. 1. 1.] -> [ True] | pred: [0.7004969  0.7199961  0.75471765] ✓
49: [1. 0. 0. 1. 1. 1. 0.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
50: [1. 0. 0. 1. 1. 0. 1.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
51: [1. 0. 0. 1. 1. 0. 0.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
52: [1. 0. 0. 1. 0. 1. 1.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
53: [1. 0. 0. 1. 0. 1. 0.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
54: [1. 0. 0. 1. 0. 0. 1.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
55: [1. 0. 0. 1. 0. 0. 0.] -> [False] | pred: [0.23323378 0.27784103 0.38287282] ✓
56: [1. 0. 0. 0. 1. 1. 1.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
57: [1. 0. 0. 0. 1. 1. 0.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
58: [1. 0. 0. 0. 1. 0. 1.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
59: [1. 0. 0. 0. 1. 0. 0.] -> [False] | pred: [0.23323378 0.27784103 0.38287282] ✓
60: [1. 0. 0. 0. 0. 1. 1.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
61: [1. 0. 0. 0. 0. 1. 0.] -> [False] | pred: [0.23323378 0.27784103 0.38287282] ✓
62: [1. 0. 0. 0. 0. 0. 1.] -> [False] | pred: [0.23323378 0.27784103 0.38287282] ✓
63: [1. 0. 0. 0. 0. 0. 0.] -> [False] | pred: [0.09584662 0.12218285 0.198593  ] ✓
64: [0. 1. 1. 1. 1. 1. 1.] -> [ True] | pred: [0.76557785 0.7768066  0.7967997 ] ✓
65: [0. 1. 1. 1. 1. 1. 0.] -> [ True] | pred: [0.7004969  0.7199961  0.75471765] ✓
66: [0. 1. 1. 1. 1. 0. 1.] -> [ True] | pred: [0.7004969  0.7199961  0.75471765] ✓
67: [0. 1. 1. 1. 1. 0. 0.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
68: [0. 1. 1. 1. 0. 1. 1.] -> [ True] | pred: [0.7004969  0.7199961  0.75471765] ✓
69: [0. 1. 1. 1. 0. 1. 0.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
70: [0. 1. 1. 1. 0. 0. 1.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
71: [0. 1. 1. 1. 0. 0. 0.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
72: [0. 1. 1. 0. 1. 1. 1.] -> [ True] | pred: [0.7004969  0.7199961  0.75471765] ✓
73: [0. 1. 1. 0. 1. 1. 0.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
74: [0. 1. 1. 0. 1. 0. 1.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
75: [0. 1. 1. 0. 1. 0. 0.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
76: [0. 1. 1. 0. 0. 1. 1.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
77: [0. 1. 1. 0. 0. 1. 0.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
78: [0. 1. 1. 0. 0. 0. 1.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
79: [0. 1. 1. 0. 0. 0. 0.] -> [False] | pred: [0.23323378 0.27784103 0.38287282] ✓
80: [0. 1. 0. 1. 1. 1. 1.] -> [ True] | pred: [0.7004969  0.7199961  0.75471765] ✓
81: [0. 1. 0. 1. 1. 1. 0.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
82: [0. 1. 0. 1. 1. 0. 1.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
83: [0. 1. 0. 1. 1. 0. 0.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
84: [0. 1. 0. 1. 0. 1. 1.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
85: [0. 1. 0. 1. 0. 1. 0.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
86: [0. 1. 0. 1. 0. 0. 1.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
87: [0. 1. 0. 1. 0. 0. 0.] -> [False] | pred: [0.23323378 0.27784103 0.38287282] ✓
88: [0. 1. 0. 0. 1. 1. 1.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
89: [0. 1. 0. 0. 1. 1. 0.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
90: [0. 1. 0. 0. 1. 0. 1.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
91: [0. 1. 0. 0. 1. 0. 0.] -> [False] | pred: [0.23323378 0.27784103 0.38287282] ✓
92: [0. 1. 0. 0. 0. 1. 1.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
93: [0. 1. 0. 0. 0. 1. 0.] -> [False] | pred: [0.23323378 0.27784103 0.38287282] ✓
94: [0. 1. 0. 0. 0. 0. 1.] -> [False] | pred: [0.23323378 0.27784103 0.38287282] ✓
95: [0. 1. 0. 0. 0. 0. 0.] -> [False] | pred: [0.09584662 0.12218285 0.198593  ] ✓
96: [0. 0. 1. 1. 1. 1. 1.] -> [ True] | pred: [0.7004969  0.7199961  0.75471765] ✓
97: [0. 0. 1. 1. 1. 1. 0.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
98: [0. 0. 1. 1. 1. 0. 1.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
99: [0. 0. 1. 1. 1. 0. 0.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
100: [0. 0. 1. 1. 0. 1. 1.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
101: [0. 0. 1. 1. 0. 1. 0.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
102: [0. 0. 1. 1. 0. 0. 1.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
103: [0. 0. 1. 1. 0. 0. 0.] -> [False] | pred: [0.23323378 0.27784103 0.38287282] ✓
104: [0. 0. 1. 0. 1. 1. 1.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
105: [0. 0. 1. 0. 1. 1. 0.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
106: [0. 0. 1. 0. 1. 0. 1.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
107: [0. 0. 1. 0. 1. 0. 0.] -> [False] | pred: [0.23323378 0.27784103 0.38287282] ✓
108: [0. 0. 1. 0. 0. 1. 1.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
109: [0. 0. 1. 0. 0. 1. 0.] -> [False] | pred: [0.23323378 0.27784103 0.38287282] ✓
110: [0. 0. 1. 0. 0. 0. 1.] -> [False] | pred: [0.23323378 0.27784103 0.38287282] ✓
111: [0. 0. 1. 0. 0. 0. 0.] -> [False] | pred: [0.09584662 0.12218285 0.198593  ] ✓
112: [0. 0. 0. 1. 1. 1. 1.] -> [ True] | pred: [0.5899636 0.6224249 0.6817076] ✓
113: [0. 0. 0. 1. 1. 1. 0.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
114: [0. 0. 0. 1. 1. 0. 1.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
115: [0. 0. 0. 1. 1. 0. 0.] -> [False] | pred: [0.23323378 0.27784103 0.38287282] ✓
116: [0. 0. 0. 1. 0. 1. 1.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
117: [0. 0. 0. 1. 0. 1. 0.] -> [False] | pred: [0.23323378 0.27784103 0.38287282] ✓
118: [0. 0. 0. 1. 0. 0. 1.] -> [False] | pred: [0.23323378 0.27784103 0.38287282] ✓
119: [0. 0. 0. 1. 0. 0. 0.] -> [False] | pred: [0.09584662 0.12218285 0.198593  ] ✓
120: [0. 0. 0. 0. 1. 1. 1.] -> [False] | pred: [0.42286056 0.4685964  0.5594897 ] ✓
121: [0. 0. 0. 0. 1. 1. 0.] -> [False] | pred: [0.23323378 0.27784103 0.38287282] ✓
122: [0. 0. 0. 0. 1. 0. 1.] -> [False] | pred: [0.23323378 0.27784103 0.38287282] ✓
123: [0. 0. 0. 0. 1. 0. 0.] -> [False] | pred: [0.09584662 0.12218285 0.198593  ] ✓
124: [0. 0. 0. 0. 0. 1. 1.] -> [False] | pred: [0.23323378 0.27784103 0.38287282] ✓
125: [0. 0. 0. 0. 0. 1. 0.] -> [False] | pred: [0.09584662 0.12218285 0.198593  ] ✓
126: [0. 0. 0. 0. 0. 0. 1.] -> [False] | pred: [0.09584662 0.12218285 0.198593  ] ✓
127: [0. 0. 0. 0. 0. 0. 0.] -> [False] | pred: [0.03262245 0.04289892 0.07766546] ✓
Accuracy: 100.0%
[10]: