Abstract

N EURAL TANGENTS is a library for working with infinite-width neural networks. It provides a high-level API for specifying complex and hierarchical neural network architectures. These networks can then be trained and evaluated either at finitewidth as usual, or in their infinite-width limit. For the infinite-width networks, N EURAL TANGENTS performs exact inference either via Bayes’ rule or gradient descent, and generates the corresponding Neural Network Gaussian Process and Neural Tangent kernels. Additionally, N EURAL TANGENTS provides tools to study gradient descent training dynamics of wide but finite networks. The entire library runs out-of-the-box on CPU, GPU, or TPU. All computations can be automatically distributed over multiple accelerators with near-linear scaling in the number of devices. N EURAL TANGENTS is available at www.github.com/neural-tangents/neural-tangents We also provide an accompanying interactive Colab notebook.1