Index

_

__init__() (neural_tangents.predict.Gaussian method) (neural_tangents.predict.ODEState method) (neural_tangents.stax.Bool method) (neural_tangents.stax.Diagonal method) (neural_tangents.stax.Padding method) (neural_tangents.stax.PositionalEmbedding method)

A

ABRelu() (in module neural_tangents.stax) Abs() (in module neural_tangents.stax)

Aggregate() (in module neural_tangents.stax) AvgPool() (in module neural_tangents.stax)

B

Bool (class in neural_tangents.stax)

C

Conv() (in module neural_tangents.stax) ConvLocal() (in module neural_tangents.stax)

ConvTranspose() (in module neural_tangents.stax) Cos() (in module neural_tangents.stax)

D

Dense() (in module neural_tangents.stax) Diagonal (class in neural_tangents.stax)

DotGeneral() (in module neural_tangents.stax) Dropout() (in module neural_tangents.stax)

E

Elementwise() (in module neural_tangents.stax) ElementwiseNumerical() (in module neural_tangents.stax) empirical_kernel_fn() (in module neural_tangents) empirical_nngp_fn() (in module neural_tangents)

empirical_ntk_fn() (in module neural_tangents) Erf() (in module neural_tangents.stax) Exp() (in module neural_tangents.stax) ExpNormalized() (in module neural_tangents.stax)

F

FanInConcat() (in module neural_tangents.stax) FanInProd() (in module neural_tangents.stax)

FanInSum() (in module neural_tangents.stax) FanOut() (in module neural_tangents.stax) Flatten() (in module neural_tangents.stax)

G

Gaussian (class in neural_tangents.predict) Gaussian() (in module neural_tangents.stax) Gelu() (in module neural_tangents.stax) GlobalAvgPool() (in module neural_tangents.stax) GlobalSelfAttention() (in module neural_tangents.stax)

GlobalSumPool() (in module neural_tangents.stax) gp_inference() (in module neural_tangents.predict) gradient_descent() (in module neural_tangents.predict) gradient_descent_mse() (in module neural_tangents.predict) gradient_descent_mse_ensemble() (in module neural_tangents.predict)

H

Hermite() (in module neural_tangents.stax)

I

Identity() (in module neural_tangents.stax)

ImageResize() (in module neural_tangents.stax) input (neural_tangents.stax.Diagonal attribute)

L

layer() (in module neural_tangents.stax) LayerNorm() (in module neural_tangents.stax)

LeakyRelu() (in module neural_tangents.stax) linearize() (in module neural_tangents)

M

max_learning_rate() (in module neural_tangents.predict) module neural_tangents neural_tangents._src.batching neural_tangents._src.empirical neural_tangents._src.monte_carlo neural_tangents._src.predict neural_tangents.stax

monte_carlo_kernel_fn() (in module neural_tangents)

N

neural_tangents module neural_tangents._src.batching module neural_tangents._src.empirical module

neural_tangents._src.monte_carlo module neural_tangents._src.predict module neural_tangents.stax module

O

ODEState (class in neural_tangents.predict)

output (neural_tangents.stax.Diagonal attribute)

P

Padding (class in neural_tangents.stax)

parallel() (in module neural_tangents.stax) PositionalEmbedding (class in neural_tangents.stax)

R

Rbf() (in module neural_tangents.stax)

Relu() (in module neural_tangents.stax) requires() (in module neural_tangents.stax)

S

serial() (in module neural_tangents.stax) Sigmoid_like() (in module neural_tangents.stax) Sign() (in module neural_tangents.stax)

Sin() (in module neural_tangents.stax) SumPool() (in module neural_tangents.stax) supports_masking() (in module neural_tangents.stax)

T

taylor_expand() (in module neural_tangents)