neural_tangents/stax.py
def kernel_fn(kernels):
    """Compute the transformed kernels after a dense layer."""
    var1, nngp, var2, ntk = \
        kernels.var1, kernels.nngp, kernels.var2, kernels.ntk

    def fc(x):
      return _affine(x, W_std, b_std)

    if parameterization == 'ntk':
      var1, nngp, var2 = map(fc, (var1, nngp, var2))
      if ntk is not None:
        ntk = nngp + W_std**2 * ntk
    elif parameterization == 'standard':
      input_width = kernels.shape1[1]
      if ntk is not None:
        ntk = input_width * nngp + 1. + W_std**2 * ntk
      var1, nngp, var2 = map(fc, (var1, nngp, var2))

    return kernels._replace(
        var1=var1, nngp=nngp, var2=var2, ntk=ntk, is_gaussian=True)
Compute the transformed kernels after a dense layer.
Compute the transformed kernels after a dense layer.
Compute the transformed kernels after a dense layer .
def kernel_fn(kernels):
    var1, nngp, var2, ntk = \
        kernels.var1, kernels.nngp, kernels.var2, kernels.ntk
    def fc(x):
      return _affine(x, W_std, b_std)
    if parameterization == 'ntk':
      var1, nngp, var2 = map(fc, (var1, nngp, var2))
      if ntk is not None:
        ntk = nngp + W_std**2 * ntk
    elif parameterization == 'standard':
      input_width = kernels.shape1[1]
      if ntk is not None:
        ntk = input_width * nngp + 1. + W_std**2 * ntk
      var1, nngp, var2 = map(fc, (var1, nngp, var2))
    return kernels._replace(
        var1=var1, nngp=nngp, var2=var2, ntk=ntk, is_gaussian=True)
def kernel_fn ( kernels ) : var1 , nngp , var2 , ntk = kernels . var1 , kernels . nngp , kernels . var2 , kernels . ntk def fc ( x ) : return _affine ( x , W_std , b_std ) if parameterization == 'ntk' : var1 , nngp , var2 = map ( fc , ( var1 , nngp , var2 ) ) if ntk is not None : ntk = nngp + W_std ** 2 * ntk elif parameterization == 'standard' : input_width = kernels . shape1 [ 1 ] if ntk is not None : ntk = input_width * nngp + 1. + W_std ** 2 * ntk var1 , nngp , var2 = map ( fc , ( var1 , nngp , var2 ) ) return kernels . _replace ( var1 = var1 , nngp = nngp , var2 = var2 , ntk = ntk , is_gaussian = True )
Compute the transformed kernels after a dense layer.
Compute the transformed kernels after a dense layer .
"""Compute the transformed kernels after a dense layer."""
kernels