<jupyter_start><jupyter_text>Optimize<jupyter_code>spec = vt.Spec('log')
def f(x, α):
    return np.log(x) - x/α
def fx(y, α):
    x = spec.decode(y)
    return f(x, α)
dfx = jax.grad(fx)
ddfx = jax.grad(dfx)
def opt(α, x0, clip=2.0, K=5):
    dobj1 = jax.partial(dfx, α=α)
    ddobj1 = jax.partial(ddfx, α=α)
    y0 = spec.encode(x0)
    y1 = vj.optim_newton(dobj1, ddobj1, y0, clip=clip, K=K)
    x1 = spec.decode(y1)
    return x1
dopt = jax.grad(opt)
opt(1.5, 0.5) # = α
dopt(1.5, 0.5) # = 1
y = np.linspace(-3.0, 3.0, 100)
plt.plot(y, fx(y, 1.5));
slv0 = vj.optimer_diff(fx, -5.0, 5.0)
slv = lambda α: spec.decode(slv0(α))
dslv = jax.value_and_grad(slv)
jdslv = jax.jit(dslv)
%time jdslv(1.5)<jupyter_output>CPU times: user 836 µs, sys: 0 ns, total: 836 µs
Wall time: 376 µs<jupyter_text>Interpolation<jupyter_code>x = np.linspace(0, 1, 101)
y = x*(1-x)
def dummy_interp(x1):
    return vj.interp(x, y, x1)
d_dummy_interp = jax.grad(dummy_interp)
dv_dummy_interp = jax.vmap(d_dummy_interp)
y1 = dummy_interp(x)
plt.plot(x, y1);
dy1 = dv_dummy_interp(x)
plt.plot(x, dy1);
plt.ylim(-1, 1);<jupyter_output><empty_output><jupyter_text>Grid Optimize<jupyter_code>N = 10000
s_grid = np.linspace(0, 1, N)
def obj(s, c):
    return -(s-c)**2
def tot(c):
    return np.sum(obj(s_grid, c))
dtot = jax.grad(tot)
jtot = jax.jit(tot)
jdtot = jax.jit(dtot)
K = 250
speed = 0.01
c = 0.5*np.ones(N)
hist = np0.zeros(K)
for i in range(K):
    hist[i] = jtot(c)
    dc = jdtot(c)
    c += speed*dc
plt.plot(hist);
plt.plot(s_grid, c - s_grid);<jupyter_output><empty_output><jupyter_text>Smoothmax Optimize<jupyter_code>x = np.linspace(-2, 2, 100)
y = 1 - x**2
plt.plot(x, y);
def obj(c, α):
    return α*c + vj.interp(x, y, c)
plt.plot(x, obj(x, 1.));
dobj = jax.grad(obj)
ddobj = jax.grad(dobj)
vddobj = jax.vmap(ddobj, in_axes=(0, None))
plt.plot(x, vddobj(x, 1.));
opt = vj.solver_diff(dobj, -2., 2.)
opt(1.)
def obj0(x, α):
    return α*x - 0.5*x**2
opt = vj.optimer_diff(obj0, -2., 2.)
vopt = jax.vmap(opt)
vopt(np.linspace(0, 1, 10))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.19601167315175097, 'fraction_numerical': 0.05155642023346303, 'mean_word_length': 4.091584158415841, 'pattern_counts': {'":': 0, '<': 16, '<?xml version=': 0, '>': 16, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '461351', 'license': 'MIT', 'n_tokens_mistral': 1138, 'n_tokens_neox': 1063, 'n_words': 206, 'path': 'notebooks/testing.ipynb', 'repo_name': 'iamlemec/valjax', 'size': '2056'}