Introduce a JAX Linker class

[brandonwillard]

This PR introduces a JAX Linker class that compiles graphs to XLA using jax.

Local PyMC3 Testing Example

For anyone interested in trying out this functionality, first install the latest PyMC3 from GitHub (e.g. pip install --no-deps git+https://github.com/pymc-devs/pymc3). We need the latest PyMC3 because it contains a Theano import fix that allows it to run using Theano-PyMC (i.e. this project). The option --no-deps is used to prevent pip from installing the old Theano. If you leave out that option, you'll have to pip uninstall theano.

Next, check out this PR branch (e.g. git clone git@github.com:brandonwillard/Theano.git -b jax-linker) and install (e.g. pip install -r requirements.txt) Theano-PyMC.

Finally, try a small example in PyMC3:

import theano

import pymc3 as pm

# Disable C compilation by default
theano.config.cxx = ""

# This will make the JAX Linker the default
theano.config.mode = "JAX"

# Create a simple model
with pm.Model() as model:
    a = pm.Exponential("a", 1.0)
    b = pm.InverseGamma("b", 0.5, 0.5)
    x = pm.Normal("x", a, b)

# Evaluate the log-likelihood in JAX!
model.logp(model.test_point)

The output should look similar to the following:

/home/bwillard/apps/anaconda3/envs/theano-36/lib/python3.7/site-packages/jax/lax/lax.py:5905: UserWarning: Explicitly requested dtype float64 requested in array is not available, and will be truncated to dtype float32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/google/jax#current-gotchas for more.
  warnings.warn(msg.format(dtype, fun_name , truncated_dtype))
/home/bwillard/apps/anaconda3/envs/theano-36/lib/python3.7/site-packages/jax/lax/lax.py:5905: UserWarning: Explicitly requested dtype float64 requested in astype is not available, and will be truncated to dtype float32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/google/jax#current-gotchas for more.
  warnings.warn(msg.format(dtype, fun_name , truncated_dtype))
Out[37]: DeviceArray(-2.749619, dtype=float32)

[junpenglao]

Interestingly I found a small differences between the gradient from theano-jax and native jax in the code snippet example:

dlogpt = theano.grad(cost=model.logpt, wrt=model.free_RVs)
dlogp_logp_fn = theano.function(inputs=model.free_RVs, outputs=[model.logpt] + dlogpt)
dlogp_logp_fn(**model.test_point)

# [DeviceArray(-2.74961873, dtype=float64),
#  DeviceArray(0.30685282, dtype=float64),
#  DeviceArray(4.4408921e-16, dtype=float64),
#  DeviceArray(0., dtype=float64)]

fgraph = theano.gof.FunctionGraph(model.free_RVs, [model.logpt])
logp_fn_jax = jax_funcify(fgraph)[0]
dlogp_logp_fn = jax.value_and_grad(logp_fn_jax, argnums=range(len(model.free_RVs)))
dlogp_logp_fn(
    model.test_point['a_log__'],
    model.test_point['b_log__'],
    model.test_point['x'])

# (DeviceArray(-2.74961873, dtype=float64),
#  (DeviceArray(0.30685282, dtype=float64),
#   DeviceArray(2.22044605e-16, dtype=float64),
#   DeviceArray(-0., dtype=float64)))

https://gist.github.com/junpenglao/fe5e1b451c076cc7b4ca16acdd7d6472

[brandonwillard]

I'm glad they're so similar, though!

Those 1e-16 values? The JAX value appears to be half the Theano value. Perhaps there's a step in there somewhere that's reducing precision. Regardless, those are both within common effectively-zero ranges, so it's probably not an error.

[junpenglao]

that makes sense - Jax is by default float32, and to enable float64 there are a bit more steps: https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#Double-(64bit)-precision

[twiecki]

As this is basically functional (even if basic) I propose we merge this and start doing individual PRs against this.

[brandonwillard]

I just pushed some major fixes that should add functionality and coverage for all the *Subtensor* Ops (i.e. all varieties of array indexing).

[twiecki]

I can now sample with jax 0.2.0.

[junpenglao]

I am getting an error of Identity lacking nfunc_spec attr, could be fixed with adding the below to jaxify.py

@jax_funcify.register(Identity)
def jax_funcify_Identity(op):
    def identity(x):
        return x

    return identity

[twiecki]

I think I'll merge, it's a bit of a nuisance we can't PR into this.

[twiecki]

OK damn, I was hoping it was only the syntax error (which I fixed on master) but other tests are failing too :-/.