yberreby

Methodology

Code like a hacker: concisely, with self-doubt, without fluff, without repeating yourself, keeping code as orthogonal as possible.

DRY

• Repeating oneself is unacceptable.
• If your LOCs look suspiciously similar, consider a loop or a lambda.
• If you can refactor to follow a "data-driven" approach (e.g. list of dicts instead of ad-hoc code), consider doing so.
• Don't be afraid of using tiny, local abstractions.

yberreby

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'

yberreby

• Use jax.tree_util.Partial to pass partially-applied functions to JIT-compiled code
• Use static_argnames, not static_argnums, whenever possible

yberreby

# %% [markdown]
# ## Diffeomorphic T → Y morph with OT loss (ott-jax)

# %%
# %matplotlib widget

# %% 0 · Imports ---------------------------------------------------------------
import numpy as np, matplotlib.pyplot as plt
from skimage.draw import line
from skimage.measure import find_contours

yberreby

# %%
# %matplotlib widget

# %%
# --- Imports ---------------------------------------------------------------
import jax
import jax.numpy as jnp
import numpy as np
import matplotlib.pyplot as plt
from scipy.sparse import diags

yberreby

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "pyqt6",         # For matplotlib backend
#   "numpy",
#   "optax>=0.2.5",
#   "matplotlib",
#   "jax[cuda12]==0.5.2", # Change for CPU.
#   "jaxtyping>=0.3.2",

yberreby

# ---
# jupyter:
#   jupytext:
#     formats: ipynb,py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.17.2
#   kernelspec:

yberreby

import functools, time, jax, jax.numpy as jnp

jax.config.update("jax_default_matmul_precision", "tensorfloat32")
SQRT2_OVER_PI = 0.7978845608028654

# ----------------------------------------------------------------------
def gelu_fast(x):
    u = SQRT2_OVER_PI * (x + 0.044715 * x * x * x)
    return 0.5 * x * (1. + jnp.tanh(u))

yberreby

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.12"
# dependencies = [
#   "jax==0.5.0",
#   "jaxopt>=0.8.5",
#   "optax>=0.2.4",
#   "matplotlib>=3.10.1",
#   "pyqt6>=6.9.0", # for matplotlib gui
# ]

yberreby

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.12"
# dependencies = [
#   "jax==0.5.0",
#   "matplotlib>=3.10.1",
#   "pyqt6>=6.9.0", # for matplotlib gui
# ]
# ///