docs: def _build(self, images):
      """Build dilation module.
  
      Args:
        images: Tensor of shape [batch_size, height, width, depth]
          and dtype float32. Represents a set of images with an arbitrary depth.
          Note that when using the default initializer, depth must equal
          num_output_classes.
  
      Returns:
        Tensor of shape [batch_size, height, width, num_output_classes] and dtype
          float32. Represents, for each image and pixel, logits for per-class
          predictions.
  
      Raises:
        IncompatibleShapeError: If images is not rank 4.
        ValueError: If model_size is not one of 'basic' or 'large'.
      """
      num_classes = self._num_output_classes
  
      if len(images.get_shape()) != 4:
        raise base.IncompatibleShapeError(
            "'images' must have shape [batch_size, height, width, depth].")
  
      if self.WEIGHTS not in self._initializers:
        if self._model_size == self.BASIC:
          self._initializers[self.WEIGHTS] = identity_kernel_initializer
        elif self._model_size == self.LARGE:
          self._initializers[self.WEIGHTS] = noisy_identity_kernel_initializer(
              num_classes)
        else:
          raise ValueError("Unrecognized model_size: %s" % self._model_size)
  
      if self.BIASES not in self._initializers:
        self._initializers[self.BIASES] = tf.zeros_initializer()
  
      if self._model_size == self.BASIC:
        self._conv_modules = [
            self._dilated_conv_layer(num_classes, 1, True, "conv1"),
            self._dilated_conv_layer(num_classes, 1, True, "conv2"),
            self._dilated_conv_layer(num_classes, 2, True, "conv3"),
            self._dilated_conv_layer(num_classes, 4, True, "conv4"),
            self._dilated_conv_layer(num_classes, 8, True, "conv5"),
            self._dilated_conv_layer(num_classes, 16, True, "conv6"),
            self._dilated_conv_layer(num_classes, 1, True, "conv7"),
            self._dilated_conv_layer(num_classes, 1, False, "conv8"),
        ]
      elif self._model_size == self.LARGE:
        self._conv_modules = [
            self._dilated_conv_layer(2 * num_classes, 1, True, "conv1"),
            self._dilated_conv_layer(2 * num_classes, 1, True, "conv2"),
            self._dilated_conv_layer(4 * num_classes, 2, True, "conv3"),
            self._dilated_conv_layer(8 * num_classes, 4, True, "conv4"),
            self._dilated_conv_layer(16 * num_classes, 8, True, "conv5"),
            self._dilated_conv_layer(32 * num_classes, 16, True, "conv6"),
            self._dilated_conv_layer(32 * num_classes, 1, True, "conv7"),
            self._dilated_conv_layer(num_classes, 1, False, "conv8"),
        ]
      else:
        raise ValueError("Unrecognized model_size: %s" % self._model_size)
  
      dilation_mod = sequential.Sequential(self._conv_modules, name="dilation")
      return dilation_mod(images)