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.

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