docs: def identity_kernel_initializer(shape, dtype=tf.float32, partition_info=None):
    """An initializer for constructing identity convolution kernels.
  
    Constructs a convolution kernel such that applying it is the same as an
    identity operation on the input. Formally, the kernel has entry [i, j, in,
    out] = 1 if in equals out and i and j are the middle of the kernel and 0
    otherwise.
  
    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 the 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 differ.
    """
    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("Identity initializer only works for square filters.")
    if filter_width % 2 != 1:
      raise ValueError(
          "Identity initializer requires filters have odd height and width.")
    if in_channels != out_channels:
      raise ValueError(
          "in_channels must equal out_channels in order to construct per-channel"
          " identities.")
  
    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 = tf.equal(
        _range_along_dimension(2, shape), _range_along_dimension(3, shape))
    return tf.cast(tf.logical_and(is_same_channel, is_middle_pixel), dtype=dtype)