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

    if cross > M.OVER_PIXELS and not is_height_width:
      filter_shape_nngp = filter_shape[::-1]
      strides_nngp = strides[::-1]
    else:
      filter_shape_nngp = filter_shape
      strides_nngp = strides

    if cross == M.OVER_PIXELS:

      def conv_nngp_unscaled(x):
        if _is_array(x):
          x = _conv_nngp_4d(x, filter_shape_nngp, strides_nngp, padding)
        return x
    elif cross in [M.OVER_POINTS, M.NO]:

      def conv_nngp_unscaled(x):
        if _is_array(x):
          x = _conv_nngp_5or6d_double_conv(x, filter_shape_nngp,
                                           strides_nngp, padding)
        return x

      is_height_width = not is_height_width
    else:
      raise NotImplementedError(
          "Only implemented for `OVER_PIXELS`, `OVER_POINTS` and `NO`;"
          " supplied {}".format(cross))

    def conv_nngp(x):
      x = conv_nngp_unscaled(x)
      return _affine(x, W_std, b_std)

    if marginal == M.OVER_PIXELS:
      def conv_var(x):
        x = _conv_var_3d(x, filter_shape_nngp, strides_nngp, padding)
        x = _affine(x, W_std, b_std)
        return x
    elif marginal in [M.OVER_POINTS, M.NO]:
      def conv_var(x):
        if _is_array(x):
          x = _conv_nngp_5or6d_double_conv(x, filter_shape_nngp,
                                           strides_nngp, padding)
        x = _affine(x, W_std, b_std)
        return x
    else:
      raise NotImplementedError(
          "Only implemented for `OVER_PIXELS`, `OVER_POINTS` and `NO`;"
          " supplied {}".format(marginal))

    var1 = conv_var(var1)
    var2 = conv_var(var2)

    if parameterization == 'ntk':
      nngp = conv_nngp(nngp)
      ntk = conv_nngp(ntk) + nngp - b_std**2 if ntk is not None else ntk
    elif parameterization == 'standard':
      nngp_unscaled = conv_nngp_unscaled(nngp)
      if ntk is not None:
        ntk = (
            input_total_dim(kernels.shape1) * nngp_unscaled + 1. +
            W_std**2 * conv_nngp_unscaled(ntk))
      nngp = _affine(nngp_unscaled, W_std, b_std)

    return kernels._replace(
        var1=var1, nngp=nngp, var2=var2, ntk=ntk, is_gaussian=True,
        is_height_width=is_height_width, marginal=marginal, cross=cross)
Compute the transformed kernels after a conv layer.
Compute the transformed kernels after a conv layer.
Compute the transformed kernels after a conv layer .
def kernel_fn(kernels):
    var1, nngp, var2, ntk, is_height_width, marginal, cross = (
        kernels.var1, kernels.nngp, kernels.var2, kernels.ntk,
        kernels.is_height_width, kernels.marginal, kernels.cross)
    if cross > M.OVER_PIXELS and not is_height_width:
      filter_shape_nngp = filter_shape[::-1]
      strides_nngp = strides[::-1]
    else:
      filter_shape_nngp = filter_shape
      strides_nngp = strides
    if cross == M.OVER_PIXELS:
      def conv_nngp_unscaled(x):
        if _is_array(x):
          x = _conv_nngp_4d(x, filter_shape_nngp, strides_nngp, padding)
        return x
    elif cross in [M.OVER_POINTS, M.NO]:
      def conv_nngp_unscaled(x):
        if _is_array(x):
          x = _conv_nngp_5or6d_double_conv(x, filter_shape_nngp,
                                           strides_nngp, padding)
        return x
      is_height_width = not is_height_width
    else:
      raise NotImplementedError(
          "Only implemented for `OVER_PIXELS`, `OVER_POINTS` and `NO`;"
          " supplied {}".format(cross))
    def conv_nngp(x):
      x = conv_nngp_unscaled(x)
      return _affine(x, W_std, b_std)
    if marginal == M.OVER_PIXELS:
      def conv_var(x):
        x = _conv_var_3d(x, filter_shape_nngp, strides_nngp, padding)
        x = _affine(x, W_std, b_std)
        return x
    elif marginal in [M.OVER_POINTS, M.NO]:
      def conv_var(x):
        if _is_array(x):
          x = _conv_nngp_5or6d_double_conv(x, filter_shape_nngp,
                                           strides_nngp, padding)
        x = _affine(x, W_std, b_std)
        return x
    else:
      raise NotImplementedError(
          "Only implemented for `OVER_PIXELS`, `OVER_POINTS` and `NO`;"
          " supplied {}".format(marginal))
    var1 = conv_var(var1)
    var2 = conv_var(var2)
    if parameterization == 'ntk':
      nngp = conv_nngp(nngp)
      ntk = conv_nngp(ntk) + nngp - b_std**2 if ntk is not None else ntk
    elif parameterization == 'standard':
      nngp_unscaled = conv_nngp_unscaled(nngp)
      if ntk is not None:
        ntk = (
            input_total_dim(kernels.shape1) * nngp_unscaled + 1. +
            W_std**2 * conv_nngp_unscaled(ntk))
      nngp = _affine(nngp_unscaled, W_std, b_std)
    return kernels._replace(
        var1=var1, nngp=nngp, var2=var2, ntk=ntk, is_gaussian=True,
        is_height_width=is_height_width, marginal=marginal, cross=cross)
def kernel_fn ( kernels ) : var1 , nngp , var2 , ntk , is_height_width , marginal , cross = ( kernels . var1 , kernels . nngp , kernels . var2 , kernels . ntk , kernels . is_height_width , kernels . marginal , kernels . cross ) if cross > M . OVER_PIXELS and not is_height_width : filter_shape_nngp = filter_shape [ : : - 1 ] strides_nngp = strides [ : : - 1 ] else : filter_shape_nngp = filter_shape strides_nngp = strides if cross == M . OVER_PIXELS : def conv_nngp_unscaled ( x ) : if _is_array ( x ) : x = _conv_nngp_4d ( x , filter_shape_nngp , strides_nngp , padding ) return x elif cross in [ M . OVER_POINTS , M . NO ] : def conv_nngp_unscaled ( x ) : if _is_array ( x ) : x = _conv_nngp_5or6d_double_conv ( x , filter_shape_nngp , strides_nngp , padding ) return x is_height_width = not is_height_width else : raise NotImplementedError ( "Only implemented for `OVER_PIXELS`, `OVER_POINTS` and `NO`;" " supplied {}" . format ( cross ) ) def conv_nngp ( x ) : x = conv_nngp_unscaled ( x ) return _affine ( x , W_std , b_std ) if marginal == M . OVER_PIXELS : def conv_var ( x ) : x = _conv_var_3d ( x , filter_shape_nngp , strides_nngp , padding ) x = _affine ( x , W_std , b_std ) return x elif marginal in [ M . OVER_POINTS , M . NO ] : def conv_var ( x ) : if _is_array ( x ) : x = _conv_nngp_5or6d_double_conv ( x , filter_shape_nngp , strides_nngp , padding ) x = _affine ( x , W_std , b_std ) return x else : raise NotImplementedError ( "Only implemented for `OVER_PIXELS`, `OVER_POINTS` and `NO`;" " supplied {}" . format ( marginal ) ) var1 = conv_var ( var1 ) var2 = conv_var ( var2 ) if parameterization == 'ntk' : nngp = conv_nngp ( nngp ) ntk = conv_nngp ( ntk ) + nngp - b_std ** 2 if ntk is not None else ntk elif parameterization == 'standard' : nngp_unscaled = conv_nngp_unscaled ( nngp ) if ntk is not None : ntk = ( input_total_dim ( kernels . shape1 ) * nngp_unscaled + 1. + W_std ** 2 * conv_nngp_unscaled ( ntk ) ) nngp = _affine ( nngp_unscaled , W_std , b_std ) return kernels . _replace ( var1 = var1 , nngp = nngp , var2 = var2 , ntk = ntk , is_gaussian = True , is_height_width = is_height_width , marginal = marginal , cross = cross )
Compute the transformed kernels after a conv layer.
Compute the transformed kernels after a conv layer .
"""Compute the transformed kernels after a conv layer."""
kernels