mattjj

### numpy version

import numpy as onp

x = onp.zeros((10, 2))
x[3:5] = 5.
print x
# [[0. 0.]
#  [0. 0.]
#  [0. 0.]

mattjj

from functools import partial

import numpy.random as npr

import jax.numpy as np
from jax import lax
from jax import grad, pjit, papply


### set up some synthetic data

mattjj

import numpy as np

def conjugate_gradient(A, b, iters, eps):
  "Approximate A^{-1} b for positive definite A using conjugate gradient."
  v = np.zeros_like(b)           # initial approximate solution
  r = b                          # initial residual vector
  rho = rho_prev = np.dot(r, r)  # initial residual norm squared
  tol = eps * np.sqrt(rho)       # residual norm stop criterion
  for k in range(iters):
    if np.sqrt(rho) < tol: break

mattjj

import types
import autograd
from autograd.tracer import trace, toposort
from autograd.core import VJPNode, add_outgrads

### closure conversion

def closure_conversion(f):
    code, globs = f.func_code, f.func_globals
    env = tuple(c.cell_contents for c in f.func_closure or ())

mattjj

import types

def closure_conversion(f):
  code, globs, freevars = f.func_code, f.func_globals, f.func_code.co_freevars
  env = dict(zip(freevars, (c.cell_contents for c in f.func_closure)))
  make_cell = lambda val: (lambda: val).func_closure[0]  # different in PY3
  def f_maker(env):
    closure = tuple(make_cell(env[name]) for name in freevars)
    return types.FunctionType(code, globs, closure=closure)
  return f_maker, env

mattjj

import numpy as np
import numpy.random as npr
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider, RadioButtons


### mniw utility functions

def sample_mniw(param, num_samples, rng=npr):
  S, M, K, nu = param

mattjj

function [alpha, k1, k2] = em(data, alpha, k1, k2, tol)
    if ~exist('tol', 'var'); tol = 1e-5; end

    div = @(x, y) bsxfun(@rdivide, x, y);
    mul = @(x, y) bsxfun(@times, x, y);
    normalize = @(X) div(X, sum(X, 1));
    loglike = @(k1, k2) mul([k1; k2], exp(mul(-[k1; k2], data)));

    while true
        % E step

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 autograd.core import primitive
from autograd.container_types import make_tuple
from autograd.util import quick_grad_check
import autograd.numpy as np

@primitive
def _f(tupl):
    a, b = tupl
    return 2.*a + 3.*b

mattjj

import autograd.numpy as np
import autograd.numpy.random as npr
from autograd import jacobian, make_hvp
import operator as op

npr.seed(0)


## direct computation