def METHOD_NAME(a, indices, mask=None):
    """
    Gather on every axis of a SparseTensor.
    :param a: SparseTensor of rank k and nnz non-zeros.
    :param indices: rank-1 int Tensor, rows and columns to keep.
    :param mask: boolean mask corresponding to indices. Computed if not provided.
    :return gathered_a: SparseTensor of the gathered input.
    :return values_mask: bool Tensor indicating surviving values, shape [nnz].
    """
    in_size = _square_size(a.dense_shape)
    out_size = tf.size(indices)
    if mask is None:
        mask = ops.indices_to_mask(indices, in_size)
    inverse_map = _indices_to_inverse_map(indices, in_size)
    return _boolean_mask_sparse_square(
        a, mask, inverse_map=inverse_map, out_size=out_size
    )