code: def _boolean_mask_sparse_square(a, mask, inverse_map, out_size):
      """
      Apply boolean_mask to every axis of a SparseTensor.
      :param a: SparseTensor with uniform dimensions and nnz non-zeros.
      :param mask: boolean mask.
      :param inverse_map: Tensor of new indices, shape [nnz]. Computed if None.
      :out_size: number of True values in mask. Computed if None.
      :return a: SparseTensor with uniform dimensions.
      :return values_mask: bool Tensor of shape [nnz] indicating valid edges.
      """
      mask = tf.convert_to_tensor(mask)
      values_mask = tf.reduce_all(
          tf.gather(mask, a.indices, axis=0), axis=-1
      )
      dense_shape = [out_size] * a.shape.ndims
      indices = tf.boolean_mask(a.indices, values_mask)
      indices = tf.gather(inverse_map, indices)
      a = tf.SparseTensor(
          indices, tf.boolean_mask(a.values, values_mask), dense_shape
      )
      return (a, values_mask)
  
file_path: spektral/spektral/layers/ops/sparse.py
imports: import tensorflow
  from tensorflow.python.ops import gen_sparse_ops
  from .None import ops
label_desc: Use image node_id 6 for calling a global function with example usage: _boolean_mask_sparse_square(a, mask, inverse_map, out_size) and returns: a,  values_mask