Skip to content

Instantly share code, notes, and snippets.

Avatar

Matthew Johnson mattjj

View GitHub Profile
View minididax.py
from __future__ import annotations
from contextlib import contextmanager
from typing import NamedTuple, Callable, Optional, Any
import numpy as np
Array = Any
class Node(NamedTuple):
vjp: Optional[Callable]
parents: List[Node]
View jax2nnc.py
# referenced @chhillee's https://github.com/pytorch/functorch/blob/main/functorch/_src/nnc_compile.py
from typing import Callable, Dict, Any, List
from functools import partial
import numpy as np
import torch
import torch._C._te as te
from jax import core
View jax2torch.py
import torch
import torch.utils.dlpack
import jax
import jax.dlpack
# A generic mechanism for turning a JAX function into a PyTorch function.
def j2t(x_jax):
x_torch = torch.utils.dlpack.from_dlpack(jax.dlpack.to_dlpack(x_jax))
return x_torch
View typeclasses.py
from typing import Callable, TypeVar
from collections import defaultdict
def ensure_tuple(x):
return x if isinstance(x, tuple) else (x,)
def safe_zip(*args):
x, *xs = args
assert all(len(x_) == len(x) for x_ in xs)
return list(zip(*args))
View bench.sh
#!/bin/bash
set -e
current_branch=$(git branch --show-current)
base=${1:-master}
alt=${2:-${current_branch}}
bench=${3:-benchmarks/api_benchmark.py}
rest="${@:4}"
View jax_primitives.py
from jax import core
# A primitive is just a name to which we associate rules.
sincos_p = core.Primitive('sincos')
# A primitive's "bind" is how it gets applied, in a way that interacts with the
# trace/transform machinery. As a convention we wrap them in Python functions
# like this:
def sincos(x):
View hmm_em.py
from functools import partial
import jax
import jax.numpy as np
from jax.scipy.special import logsumexp
from jax import lax, random
from jax import jit, grad
def log_normalizer(params, seq):
View rewrite.py
from functools import partial
from jax import core
from jax.util import safe_map, safe_zip
import jax.linear_util as lu
map = safe_map
zip = safe_zip
View jax_custom_vjps.py
from jax.interpreters import ad
from jax.interpreters import partial_eval as pe
from jax import custom_transforms
from jax import core
from jax import grad
@custom_transforms
def f(x, y):
return x**2 + 3 * y
View jax_taylor.py
from functools import partial
from math import factorial
import jax.numpy as np
import matplotlib.pyplot as plt
from jax import jvp, vmap
def f(x):
return 1./5 * x**3 + 3 * x**2 - x + 1.