docs: def _apply_mask(self):
      """Applies the passed-in mask to the convolution matrix.
  
      Returns:
        w: A copy of the convolution matrix that has had the mask applied.
  
      Raises:
        base.IncompatibleShapeError: If the mask shape has more dimensions than
            the weight matrix.
        base.IncompatibleShapeError: If the mask and the weight matrix don't
            match on shape.
      """
      w = self._w
      w_shape = w.get_shape()
      mask_shape = self._mask.get_shape()
  
      if mask_shape.ndims > w_shape.ndims:
        raise base.IncompatibleShapeError(
            "Invalid mask shape: {}. Max shape: {}".format(
                mask_shape.ndims, len(self._data_format)
            )
        )
      if mask_shape != w_shape[:mask_shape.ndims]:
        raise base.IncompatibleShapeError(
            "Invalid mask shape: {}. Weight shape: {}".format(
                mask_shape, w_shape
            )
        )
      # TF broadcasting is a bit fragile.
      # Expand the shape of self._mask by one dim at a time to the right
      # until the rank matches `weight_shape`.
      while self._mask.get_shape().ndims < w_shape.ndims:
        self._mask = tf.expand_dims(self._mask, -1)
  
      # tf.Variable & tf.ResourceVariable don't support *=.
      w = w * self._mask  # pylint: disable=g-no-augmented-assignment
  
      return w