mattjj

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

mattjj

import numpy as np

def gram_matrix(Xs):
  temp = np.vstack([np.ravel(X) for X in Xs])
  return np.dot(temp, temp.T)

def eig(X):
  vals, vecs = np.linalg.eig(X)
  idx = np.argsort(np.abs(vals))
  return vals[idx], vecs[...,idx]

mattjj

from __future__ import division
import numpy as np
from numpy.random import randn
from scipy.linalg.blas import drot, drotg

# references for updates:
#   - Golub and van Loan (4th ed.) Section 6.5.4
#   - http://mathoverflow.net/questions/30162/is-there-a-way-to-simplify-block-cholesky-decomposition-if-you-already-have-deco
#
# references for downdates:

mattjj

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]

mattjj

--recurse=yes
--tag-relative=yes
--exclude=*.git*
--exclude=*.pyc
--exclude=*.pyo
--exclude=.DS_Store
--exclude=*.md
--exclude=*.mkd

--langdef=Clojure

mattjj

var socket = io.connect(window.location.host);
var received;
socket.on('data', function (data) {
    received = JSON.parse(data);
});

mattjj

# 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

mattjj

#!/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}"

mattjj

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.

mattjj

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))