I made a framework for JAX visualisation where you can annotate your JAX-compatible function with it, and after it runs it generates a graph with the output of each intermediate operation that JAX generated. There is also another visualizer, which visualizes the transformation that happens to the function's input when the function is applied. You can think about grad as one such transformation.
I think seeing these formulas can be very helpful in learning them, and the generated images can be another learning channel for multimodal models or for agents to "see" what equations look like.
On this website, I am gonna experiment with different functions each day and publish if they turned out alright. Btw, to generate the images, I just need to annotate them with jaxvis.draw.
cdf · gradients · outliers · distributions
import jax
import jax.numpy as jnp
import numpy as np
import jaxvis
import fad
from jaxvis.render import ramp
fad.day(__file__)
n = 80_000
rng = np.random.default_rng(0)
cloud = dict(rep={(n, 2): "points"}, only="points", structural=True,
size=640, tween=30, hold=8, duration=55, blend="mean")
def by_value(v):
u = (v - v.min()) / (np.ptp(v) + 1e-9)
c = ramp("ultra")[(u * 255).astype(int)] / 255.0
return c / c.max(1, keepdims=True)
$$u = \Phi(x) = \tfrac{1}{2}\left(1 + \operatorname{erf}(x/\sqrt{2})\right)$$
Push any distribution through its own CDF to a uniform distribution.
gauss = rng.normal(size=(n, 2))
@jaxvis.draw(jnp.asarray(gauss), palette="bloom",
colors=by_value(np.hypot(gauss[:, 0], gauss[:, 1])),
frame="fixed", **cloud)
def to_uniform(x):
return jax.scipy.special.erf(x / jnp.sqrt(2.0))
$$r = \sqrt{-2\ln u_1}, \quad \theta = 2\pi u_2, \quad z = r(\cos\theta, \sin\theta)$$
unif = rng.uniform(1e-4, 1.0, size=(n, 2))
@jaxvis.draw(jnp.asarray(unif), palette="bloom", colors=by_value(unif[:, 0]),
**cloud)
def box_muller(u):
r = jnp.sqrt(-2.0 * jnp.log(u[:, :1]))
th = 2.0 * jnp.pi * u[:, 1:]
return jnp.concatenate([jnp.cos(th), jnp.sin(th)], 1) * r
import jax
import jax.numpy as jnp
import numpy as np
import jaxvis
import fad
from jaxvis.fields import grid
fad.day(__file__)
$$z \leftarrow z + \eta\,(y - \sigma(z)) + \eta\lambda\,\nabla^2 z$$
Start from pure noise and descend on two terms at once. One pulls pixels to the edges, one punishes disagreement with neighbours.
gx, gy = grid(256, 1.0)
target = jnp.asarray((np.hypot(np.asarray(gx), np.asarray(gy)) < 0.55)
.astype(float))
noise = jnp.asarray(np.random.default_rng(0).normal(size=(256, 256)) * 2.5)
lr, smooth = 0.06, 3.0
def lap(a):
return (jnp.roll(a, 1, 0) + jnp.roll(a, -1, 0) +
jnp.roll(a, 1, 1) + jnp.roll(a, -1, 1) - 4 * a)
@jaxvis.draw_sim(init=noise, frames=130, every=1, tween=2, palette="ultra",
show=jax.nn.sigmoid, size=680, duration=45, grain=0.0)
def a_circle_from_noise(z):
return z + lr * (target - jax.nn.sigmoid(z)) + lr * smooth * lap(z)
$$z \leftarrow z + \eta\,(y - \sigma(z))$$
Drop the neighbour term.
@jaxvis.draw_sim(init=noise, frames=130, every=1, tween=2, palette="ultra",
show=jax.nn.sigmoid, size=680, duration=45, grain=0.0)
def cross_entropy(z):
return z + lr * (target - jax.nn.sigmoid(z))
$$z \leftarrow z + \eta\,(z - z^3) + \eta\lambda\,\nabla^2 z$$
Swap the label for a term that pushes each pixel away from zero, toward either $+1$ or $-1$, and keep the smoothing. Nothing is being fitted now.
@jaxvis.draw_sim(init=noise * 0.04, frames=130, every=4, tween=2,
palette="ultra", show=jax.nn.sigmoid, size=680, duration=45,
grain=0.0)
def no_target_at_all(z):
return z + lr * (z - z ** 3) + lr * smooth * lap(z)
import jax.numpy as jnp
import numpy as np
import jaxvis
import fad
fad.day(__file__)
n = 80_000
rng = np.random.default_rng(1)
cloud = dict(rep={(n, 2): "points"}, only="points", structural=True,
size=640, tween=30, hold=8, duration=55, blend="mean",
frame=(0.0, 0.0, 5.5))
clean = rng.normal(size=(n, 2))
bad = rng.integers(0, 100, n) < 3
data = jnp.asarray(np.where(bad[:, None], clean * 9.0, clean))
flag = np.where(bad[:, None], [1.0, 0.35, 0.5], [0.3, 0.8, 1.0])
$$z = \frac{x - \bar{x}}{s}$$
Standardising with the mean and standard deviation.
mean = np.asarray(data).mean(0)
std = np.asarray(data).std(0)
@jaxvis.draw(data, palette="bloom", colors=flag, **cloud)
def z_score(x):
return (x - mean) / std
$$z = \frac{x - \mathrm{med}(x)}{1.4826\,\mathrm{MAD}}$$
med = np.median(np.asarray(data), 0)
mad = 1.4826 * np.median(np.abs(np.asarray(data) - med), 0)
@jaxvis.draw(data, palette="bloom", colors=flag, **cloud)
def robust_scale(x):
return (x - med) / mad
$$x \mapsto \mathrm{clip}(x,\; -3,\; 3)$$
Rather than dropping outliers, pin them to the edge.
@jaxvis.draw(data, palette="bloom", colors=flag, **cloud)
def winsorise(x):
return jnp.clip((x - med) / mad, -3.0, 3.0)
import jax
import jax.numpy as jnp
import numpy as np
import jaxvis
import fad
fad.day(__file__)
n = 80_000
rng = np.random.default_rng(0)
cloud = dict(rep={(n, 2): "points"}, only="points", structural=True,
size=640, tween=30, hold=8, duration=55)
$$S_k = \frac{R + \sum_{i=1}^{k} U_i}{\sqrt{1 + k\,\sigma_U^2}}$$
Simulate the central limit theorem by starting with any distribution and adding plain uniform noise.
steps, amp = 200, 0.275
su = amp / np.sqrt(3.0)
dirs = np.array([[0.0, 1.0], [-0.866, -0.5], [0.866, -0.5]]) * 2.6
squares = jnp.asarray(dirs[rng.integers(0, 3, n)]
+ rng.uniform(-0.85, 0.85, size=(n, 2)))
key = jax.random.PRNGKey(1)
@jaxvis.draw_sim(init=(squares, jnp.asarray(0)), frames=steps, every=1,
tween=1, rep="points", palette="vapor", size=640,
duration=45,
show=lambda s: s[0] / jnp.sqrt(1.0 + s[1] * su ** 2))
def central_limit(state):
total, k = state
draw = jax.random.uniform(jax.random.fold_in(key, k), total.shape,
minval=-amp, maxval=amp)
return total + draw, k + 1
$$z = L^{-1}(x - \mu), \qquad LL^{\top} = \Sigma$$
tilted = rng.multivariate_normal([0.9, -0.5], [[2.2, 1.6], [1.6, 1.4]], n)
mu = tilted.mean(0)
chol = np.linalg.cholesky(np.cov(tilted.T))
unmix = np.linalg.inv(chol).T
ang = (np.arctan2(tilted[:, 1] - mu[1], tilted[:, 0] - mu[0]) + np.pi) / (2 * np.pi)
wheel = np.stack([np.sin(np.pi * ang), np.sin(np.pi * (ang + 1 / 3)),
np.sin(np.pi * (ang + 2 / 3))], 1) ** 2
wheel = wheel / wheel.max(1, keepdims=True)
@jaxvis.draw(jnp.asarray(tilted), palette="bloom", colors=wheel, blend="mean",
frame="fixed", **cloud)
def whiten(x):
return (x - mu) @ unmix