Index

_

__init__() (neural_tangents.predict.Gaussian method) (neural_tangents.predict.ODEState method) (neural_tangents.stax.AggregateImplementation method) (neural_tangents.stax.AttentionMechanism 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 (neural_tangents.stax.AttentionMechanism attribute) Abs() (in module neural_tangents.stax) Aggregate() (in module neural_tangents.stax) AggregateImplementation (class in neural_tangents.stax) AnalyticKernelFn (class in neural_tangents._src.utils.typing)

ApplyFn (class in neural_tangents._src.utils.typing) asdict() (neural_tangents.Kernel method) astuple() (neural_tangents.Kernel method) AttentionMechanism (class in neural_tangents.stax) AUTO (neural_tangents.NtkImplementation attribute) AvgPool() (in module neural_tangents.stax) Axes (in module neural_tangents._src.utils.typing)

B

batch_axis (neural_tangents.Kernel attribute)

Bool (class in neural_tangents.stax)

C

channel_axis (neural_tangents.Kernel attribute) CIRCULAR (neural_tangents.stax.Padding attribute) CONCAT (neural_tangents.stax.PositionalEmbedding attribute) Conv() (in module neural_tangents.stax) ConvLocal() (in module neural_tangents.stax)

ConvTranspose() (in module neural_tangents.stax) Cos() (in module neural_tangents.stax) cov1 (neural_tangents.Kernel attribute) cov2 (neural_tangents.Kernel attribute) covariance (neural_tangents.predict.Gaussian attribute)

D

DENSE (neural_tangents.stax.AggregateImplementation attribute) Dense() (in module neural_tangents.stax) Diagonal (class in neural_tangents.stax) diagonal_batch (neural_tangents.Kernel attribute)

diagonal_spatial (neural_tangents.Kernel attribute) dot_general() (neural_tangents.Kernel method) 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) empirical_ntk_fn_tf() (in module neural_tangents.experimental)

empirical_ntk_vp_fn() (in module neural_tangents) EmpiricalGetKernelFn (class in neural_tangents._src.utils.typing) EmpiricalKernelFn (class in neural_tangents._src.utils.typing) 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) fx_test (neural_tangents.predict.ODEState attribute) fx_train (neural_tangents.predict.ODEState attribute)

G

Gabor() (in module neural_tangents.stax) Gaussian (class in neural_tangents.predict) Gaussian() (in module neural_tangents.stax) Gelu() (in module neural_tangents.stax) get_apply_fn_and_params() (in module neural_tangents.experimental) 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 (neural_tangents.stax.AttentionMechanism attribute) Identity() (in module neural_tangents.stax) ImageResize() (in module neural_tangents.stax) Index() (in module neural_tangents.stax)

InitFn (class in neural_tangents._src.utils.typing) input (neural_tangents.stax.Diagonal attribute) is_gaussian (neural_tangents.Kernel attribute) is_input (neural_tangents.Kernel attribute) is_reversed (neural_tangents.Kernel attribute)

J

JACOBIAN_CONTRACTION (neural_tangents.NtkImplementation attribute)

K

Kernel (class in neural_tangents)

Kernels (in module neural_tangents._src.utils.typing)

L

layer() (in module neural_tangents.stax) LayerKernelFn (class in neural_tangents._src.utils.typing)

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

M

mask() (neural_tangents.Kernel method) mask1 (neural_tangents.Kernel attribute) mask2 (neural_tangents.Kernel attribute) MaskFn (class in neural_tangents._src.utils.typing) max_learning_rate() (in module neural_tangents.predict) MAYBE (neural_tangents.stax.Bool attribute) mean (neural_tangents.predict.Gaussian attribute) module neural_tangents neural_tangents._src.batching neural_tangents._src.empirical neural_tangents._src.monte_carlo neural_tangents._src.predict neural_tangents._src.utils.typing neural_tangents.experimental neural_tangents.stax

Monomial() (in module neural_tangents.stax) monte_carlo_kernel_fn() (in module neural_tangents) MonteCarloKernelFn (class in neural_tangents._src.utils.typing)

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._src.utils.typing module

neural_tangents.experimental module neural_tangents.stax module nngp (neural_tangents.Kernel attribute) NO (neural_tangents.stax.Bool attribute) NONE (neural_tangents.stax.PositionalEmbedding attribute) ntk (neural_tangents.Kernel attribute) NTK_VECTOR_PRODUCTS (neural_tangents.NtkImplementation attribute) NtkImplementation (class in neural_tangents) NTTree (in module neural_tangents._src.utils.typing) NTTrees (in module neural_tangents._src.utils.typing)

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)

Polynomial() (in module neural_tangents.stax) PositionalEmbedding (class in neural_tangents.stax) PyTree (in module neural_tangents._src.utils.typing)

Q

qx_test (neural_tangents.predict.ODEState attribute)

qx_train (neural_tangents.predict.ODEState attribute)

R

Rbf() (in module neural_tangents.stax) RectifiedMonomial() (in module neural_tangents.stax) RELU (neural_tangents.stax.AttentionMechanism attribute) Relu() (in module neural_tangents.stax)

repeat() (in module neural_tangents.stax) replace() (neural_tangents.Kernel method) requires() (in module neural_tangents.stax) reverse() (neural_tangents.Kernel method)

S

SAME (neural_tangents.stax.Padding attribute) serial() (in module neural_tangents.stax) shape1 (neural_tangents.Kernel attribute) shape2 (neural_tangents.Kernel attribute) Shapes (in module neural_tangents._src.utils.typing) Sigmoid_like() (in module neural_tangents.stax) Sign() (in module neural_tangents.stax)

Sin() (in module neural_tangents.stax) Slice (in module neural_tangents.stax) SOFTMAX (neural_tangents.stax.AttentionMechanism attribute) SPARSE (neural_tangents.stax.AggregateImplementation attribute) STRUCTURED_DERIVATIVES (neural_tangents.NtkImplementation attribute) SUM (neural_tangents.stax.PositionalEmbedding attribute) SumPool() (in module neural_tangents.stax) supports_masking() (in module neural_tangents.stax)

T

taylor_expand() (in module neural_tangents)

transpose() (neural_tangents.Kernel method)

V

VALID (neural_tangents.stax.Padding attribute)

VMapAxes (in module neural_tangents._src.utils.typing)

X

x1_is_x2 (neural_tangents.Kernel attribute)

Y

YES (neural_tangents.stax.Bool attribute)