def METHOD_NAME(self, a, s, **kwargs):
    a_pool = ops.matmul_at_b_a(s, a)
    # Modularity loss
    mod_loss = self.modularity_loss(a, s, a_pool)
    if K.ndim(a) == 3:
        mod_loss = K.mean(mod_loss)
    self.add_loss(mod_loss)
    return a_pool