Code

Fourier Sliced-Wasserstein (FSW) Embedding

Efficient implementation of the FSW embedding—a bounded-distortion Euclidean embedding for multisets over $\mathbb{R}^d$. It serves as a drop-in replacement for sum-pooling, which provably incurs infinite distortion.

• Supports weighted multisets (maintains the bounded-distortion guarantee assuming the weights come from a finite set).

Part of our paper titled “Fourier Sliced-Wasserstein Embedding for Multisets and Measures”
PyPI GitHub

Install the PyTorch package:

pip install fswlib

Fourier Sliced-Wasserstein Graph Neural Network (FSW-GNN)

A graph neural network with bounded-distortion 1-WL separation power for graphs with vertex and edge features.

• When randomly initialized, this GNN computes a Euclidean embedding for graphs that is bi-Lipschitz with respect to two standard graph metrics.
• Supports continuous edge weights while maintaining 1-WL separation (but not bi-Lipschitzness, which is provably impossible).

Part of our paper titled “FSW-GNN: A Bi-Lipschitz WL-Equivalent Graph Neural Network”
PyTorch Geometric

Sparse Approximation by the Generalized Soft-Min Penalty

Solver for the sparse approximation problem, based on the Generalized Soft-Min (GSM) penalty. Achieved state-of-the-art results at the time of publication.

Part of our paper titled “The Trimmed Lasso: Sparse Recovery Guarantees and Practical Optimization by the Generalized Soft-Min Penalty”
Matlab, C Python

Install the Python package:

pip install sparse-approx-gsm

Python implementation kindly contributed by Shachar Cohen (Weizmann).