Build an initializer for constructing near-identity convolution kernels.

  Construct a convolution kernel where in_channels and out_channels are
  multiples of base_num_channels, but need not be equal. This initializer is
  essentially the same as identity_kernel_initializer, except that magnitude
  is "spread out" across multiple copies of the input.

  Args:
    base_num_channels: int. Number that divides both in_channels and
      out_channels.
    stddev: float. Standard deviation of truncated normal noise added to
      off-entries to break ties.

  Returns:
    Initializer function for building a noisy identity kernel.

def noisy_identity_kernel_initializer(base_num_channels, stddev=1e-8):
  """Build an initializer for constructing near-identity convolution kernels.

  Construct a convolution kernel where in_channels and out_channels are
  multiples of base_num_channels, but need not be equal. This initializer is
  essentially the same as identity_kernel_initializer, except that magnitude
  is "spread out" across multiple copies of the input.

  Args:
    base_num_channels: int. Number that divides both in_channels and
      out_channels.
    stddev: float. Standard deviation of truncated normal noise added to
      off-entries to break ties.

  Returns:
    Initializer function for building a noisy identity kernel.
  """

  # pylint: disable=unused-argument
  def _noisy_identity_kernel_initializer(shape,
                                         dtype=tf.float32,
                                         partition_info=None):
    """Constructs a noisy identity kernel.

    Args:
      shape: List of integers. Represents shape of result.
      dtype: data type for values in result.
      partition_info: Partition information for initializer functions. Ignored.

    Returns:
      Tensor of desired shape and dtype such that applying it as a convolution
        kernel results in a noisy near-identity operation.

    Raises:
      ValueError: If shape does not define a valid kernel.
                  If filter width and height differ.
                  If filter width and height are not odd numbers.
                  If number of input and output channels are not multiples of
                    base_num_channels.
    """
    if len(shape) != 4:
      raise ValueError("Convolution kernels must be rank 4.")

    filter_height, filter_width, in_channels, out_channels = shape

    if filter_width != filter_height:
      raise ValueError(
          "Noisy identity initializer only works for square filters.")
    if filter_width % 2 != 1:
      raise ValueError(
          "Noisy identity initializer requires filters have odd height and "
          "width.")
    if (in_channels % base_num_channels != 0 or
        out_channels % base_num_channels != 0):
      raise ValueError("in_channels and out_channels must both be multiples of "
                       "base_num_channels.")

    middle_pixel = filter_height // 2
    is_middle_pixel = tf.logical_and(
        tf.equal(_range_along_dimension(0, shape), middle_pixel),
        tf.equal(_range_along_dimension(1, shape), middle_pixel))
    is_same_channel_multiple = tf.equal(
        tf.floordiv(
            _range_along_dimension(2, shape) * base_num_channels, in_channels),
        tf.floordiv(
            _range_along_dimension(3, shape) * base_num_channels, out_channels))
    noise = tf.truncated_normal(shape, stddev=stddev, dtype=dtype)
    return tf.where(
        tf.logical_and(is_same_channel_multiple, is_middle_pixel),
        tf.ones(
            shape, dtype=dtype) * (base_num_channels / out_channels),
        noise)

  return _noisy_identity_kernel_initializer