Jax

Numpy + Autograd. Use XLA to compile and run NumPy code on accelerators. Asynchronous dispatch, for sync use block_until_ready() import jax.numpy as jnp from jax import random key = random.PRNGKey(0) x = random.normal(key, (10,)) jnp.dot(x, x.T).block_until_ready() notable function: jit() for compilation of multiple computations grad() for performing transformation (autodiff, Jacobian-vector product) vmap() for auto-vectorisation Arrays are immutable in Jax Treat functions as pure as to compiled with XLA entropix/dslider.pyfrom functools import partial from typing import NamedTuple, Tuple import jax import jax.numpy as jnp import jax.scipy as jsp @jax.jit def kl_divergence(logp: jnp.ndarray, logq: jnp.ndarray) -> jnp.ndarray: """Compute KL divergence between two log probability distributions.""" p = jnp.exp(logp) return jnp.sum(jnp.where(p > 0, p * (logp - logq), 0.0), axis=-1) @jax.jit def ent_varent(logp: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]: """Compute entropy and varentropy from log probabilities.""" p = jnp.exp(logp) ent = -jnp.sum(p * logp, axis=-1) diff = logp + ent[..., None] varent = jnp.sum(p * diff**2, axis=-1) return ent, varent @jax.jit def normalize_logits(logits: jnp.ndarray, noise_floor: float) -> jnp.ndarray: """Normalize logits to log probabilities with noise floor truncation.""" shifted = logits - jnp.max(logits, axis=-1, keepdims=True) normalized = shifted - jax.nn.logsumexp(shifted + EPS, axis=-1, keepdims=True) # noise floor calculated for bfloat16 return jnp.where(normalized < noise_floor, jnp.log(EPS), normalized) references: github control flow see also link The following works: @jax.jit def f(x): for i in range(3): x = 2 * x return x print(f(3)) @jax.jit def g(x): y = 0.