code: def boolean_mask_sparse_square(a, mask):
      """
      Apply mask to every axis of SparseTensor a.
      :param a: SparseTensor, square, nnz non-zeros.
      :param mask: boolean mask with size equal to each dimension of a.
      :return masked_a: SparseTensor
      :return values_mask: bool tensor of shape [nnz] indicating valid values.
      """
      i = tf.squeeze(tf.where(mask), axis=-1)
      out_size = tf.size(i)
      in_size = _square_size(a.dense_shape)
      inverse_map = _indices_to_inverse_map(i, in_size)
      return _boolean_mask_sparse_square(
          a, mask, inverse_map=inverse_map, out_size=out_size
      )
  
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 8 for calling a global function with example usage: boolean_mask_sparse_square(a, mask) and returns: _boolean_mask_sparse_square