neural_tangents.stax.Rbf

neural_tangents.stax.Rbf(gamma=1.0)

Dual activation function for normalized RBF or squared exponential kernel.

Dual activation function is f(x) = sqrt(2)*sin(sqrt(2*gamma) x + pi/4). NNGP kernel transformation correspond to (with input dimension d) k = exp(- gamma / d * ||x - x'||^2) = exp(- gamma*(q11 + q22 - 2 * q12)).

Parameters:

gamma (float) – related to characteristic length-scale (l) that controls width of the kernel, where gamma = 1 / (2 l^2).

Return type:

tuple[InitFn, ApplyFn, LayerKernelFn]

Returns:

(init_fn, apply_fn, kernel_fn).