This repository contains a number of subprojects related to the interlinking of computation and dynamics in artificial and biological neural systems.
This is not an officially supported Google product.
The code is written in Python 2.7.13. You will also need:
- JAX version 1.1.18 (install) -
- NumPy, SciPy, Matplotlib (install SciPy stack, contains all of them)
- h5py (install)
- A GPU - XLA compiles these examples to CPU very slowly, so best to use a GPU for now.
LFADS is a tool for inferring dynamics from noisy, high-dimensional observations of a dynamics system. It is a sequential auto-encoder with some very particular bells and whistles. Here we have released a tutorial version, written in Python / Numpy / JAX intentionally implemented with readabilily, comprehension and innovation in mind. You may find the full TensorFlow implementation with run manager support (here).
The LFADS tutorial uses the integrator RNN example (see below). The LFADS tutorial example attempts to infer the hidden states of the integrator RNN as well as the white noise input to the RNN. One runs the integrator RNN example and then copies the resulting data file, written in /tmp/ to /tmp/LFADS/data/. Edit the name of the data file in run_lfads.py and then run execute run_lfads.py.
The LFADS tutorial is run through this Jupyter notebook.
Integration is a very simple task and highlights how to set up a loop over time, batch over multiple input/target examples, use just-in-time compilation to speed the computation up, and take a gradient in JAX. The data from this example is also used as input for the LFADS tutorial.
This example is run through this Jupyter notebook.
In Colab, Train an echostate network (ESN) to generate the chaotic output of another recurrent neural network. This Colab / IPython notebook implements a continuous-time ESN with FORCE learning implemented via recursive least squares (RLS). It also lets you use a GPU and quickly get started with JAX! Two different implementations are explored, one at the JAX / Python level and another at the LAX level. After JIT compilation, the JAX implementation runs very fast.