floydnet 0.1.0__tar.gz → 0.1.1__tar.gz

{floydnet-0.1.0 → floydnet-0.1.1}/.gitignore

@@ -205,3 +205,14 @@ cython_debug/
 marimo/_static/
 marimo/_lsp/
 __marimo__/
+
+example/data/count/processed/
+example/data/count/hom.npy
+example/data/count/iso.npy
+example/wandb
+example/output
+example/outputs
+.github/
+count.out
+run_scripts
+example/data/TSP

floydnet-0.1.1/CHANGELOG.md

@@ -0,0 +1,12 @@
+# Changelog
+
+All notable changes to this project will be documented in this file.
+
+## [1.0.0] - 2026-01-25
+- Full release with training and evaluation scripts for Graph Count, BREC, and TSP.
+- Added `pivotal_attention3` functional API for 3-Floyd attention.
+- Added additional configuration options in `PivotalAttentionBlock`.
+
+## [0.1.0] - 2025-10-21
+- Initial public skeleton with module + functional attention
+- examples scaffolding

{floydnet-0.1.0 → floydnet-0.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: floydnet
-Version: 0.1.0
+Version: 0.1.1
 Summary: Floyd Multi-Head Attention: a drop-in variant of PyTorch MHA with module and function APIs
 Project-URL: Homepage, https://github.com/ocx-lab/FloydNet
 Project-URL: Repository, https://github.com/ocx-lab/FloydNet
@@ -230,49 +230,162 @@ Requires-Dist: pytest>=7.4; extra == 'dev'
 Requires-Dist: ruff>=0.5; extra == 'dev'
 Description-Content-Type: text/markdown
 
-# floyd-net
+# FloydNet
 
-Floyd Multi-Head Attention (F-MHA) is a drop-in variant of PyTorch's attention stack. It provides:
+Official implementation of an ICLR paper (TODO: add paper title, authors, and links/arXiv).
 
-- Module API: `FloydMultiheadAttention` mirroring `torch.nn.MultiheadAttention`
-- Functional API: `floyd_scaled_dot_product_attention` mirroring `torch.nn.functional.scaled_dot_product_attention`
+![Figure Pivotal Attention Mechanism for 2-Floyd/3-Floyd.](misc/pivotalattn2&3.png)
 
-Install and manage with `uv` for a modern Python workflow.
+This repository serves two audiences:
 
-## Quick start
+- **Engineering users**: reusable PyTorch components (functional attention APIs and Transformer-style blocks) under `src/`.
+- **Research users**: scripts/configs to reproduce paper experiments under `example/`.
 
+---
+
+## Introduction
+
+FloydNet is the official PyTorch implementation accompanying an ICLR paper (TODO).  
+The repository provides:
+
+1. **Reusable components**: a drop-in attention/Transformer-block interface intended for integration into existing projects.
+2. **Reproduction code**: end-to-end training/evaluation pipelines to reproduce the benchmarks reported in the paper.
+
+For algorithmic details, hyperparameter choices, and analysis, please refer to the paper (TODO: link).
+
+---
+
+## Repository Structure
+
+- `src/floydnet/`  
+  **Library code for reuse**  
+  Contains the functional attention API and module/block implementations.
+
+- `example/`  
+  **Experiment reproduction code**  
+  Includes benchmark-specific scripts, configs, and data preparation utilities.
+
+---
+
+## Using the Attention / Transformer Block
+
+This section targets **engineering users** who want to import FloydNet as a dependency.
+
+### Installation
+
+#### Option A: Install from PyPI
 ```bash
-# Install with uv (recommended)
-uv venv --python 3.10
-source .venv/bin/activate
-uv pip install -e .[dev]
+pip install floydnet
 ```
 
-```python
-import torch
-from floyd_net import FloydMultiheadAttention
+#### Option B: Install from source
+```bash
+git clone git@github.com:ocx-lab/FloydNet.git
+cd FloydNet
+pip install -e .
+```
 
-m = FloydMultiheadAttention(embed_dim=64, num_heads=8, batch_first=True)
-x = torch.randn(2, 16, 64)
-out, attn = m(x, x, x)
-print(out.shape)
+> Requirements: Python `>= 3.9`, PyTorch `>= 2.1` (see `pyproject.toml`).
+
+### Public API
+
+FloydNet re-exports the public API from `src/floydnet/__init__.py`, so you can import from the top-level package:
+
+- **Functional API**:
+  - `pivotal_attention` (see `src/floydnet/functional.py`)
+- **Module / block API**:
+  - `PivotalAttentionBlock` (see `src/floydnet/transformer.py`)
+
+```python
+from floydnet import pivotal_attention, PivotalAttentionBlock
 ```
 
-### Functional API
+### Minimal usage example
+
 ```python
 import torch
-import torch.nn.functional as F
-from floyd_net import floyd_scaled_dot_product_attention
+from floydnet import pivotal_attention, PivotalAttentionBlock
+
+# -------------------------
+# Module API (Transformer-style block)
+# Input is a 2D grid: (B, N, N, C)
+# -------------------------
+B, N, C = 2, 16, 64
+x = torch.randn(B, N, N, C)
 
-q = torch.randn(2, 8, 16, 64)  # (B, H, L, D)
-k = torch.randn(2, 8, 16, 64)
-v = torch.randn(2, 8, 16, 64)
-out = floyd_scaled_dot_product_attention(q, k, v)
+m = PivotalAttentionBlock(embed_dim=C, num_heads=8, dropout=0.0)
+out = m(x)  # (B, N, N, C)
 print(out.shape)
+
+# -------------------------
+# Functional API
+# All inputs are 5D: (B, H, N, N, D)
+# -------------------------
+B, H, N, D = 2, 8, 16, 64
+q_ik = torch.randn(B, H, N, N, D)
+k_ij = torch.randn(B, H, N, N, D)
+k_jk = torch.randn(B, H, N, N, D)
+v_ij = torch.randn(B, H, N, N, D)
+v_jk = torch.randn(B, H, N, N, D)
+
+y = pivotal_attention(q_ik, k_ij, k_jk, v_ij, v_jk)  # (B, H, N, N, D)
+print(y.shape)
 ```
 
-## Paper reproductions
-See `paper/` for dataset preparation, configs, and experiment templates to reproduce the results in the paper.
+---
+
+## Reproducing Paper Results
+
+This section targets **research users** who want to reproduce the experiments in the paper.
+
+See `example/README.md` For detailed description.
+
+### Environment setup
+
+We recommend using `uv` to create an isolated environment for the reproduction code under `example/`.
+
+```bash
+cd /path/to/FloydNet
+
+# 1) Create a uv virtual environment with Python 3.12
+uv venv --python 3.12
+
+# 2) Activate it
+source .venv/bin/activate
+
+# 3) Install extra dependencies for reproducing paper experiments
+uv pip install -r example/requirements.txt
+
+# 4) Install FloydNet (editable) for local development / imports
+uv pip install -e .
+```
+
+## Changelog (latest)
+
+- Full release with training and evaluation scripts for Graph Count, BREC, and TSP.
+- Added `pivotal_attention3` functional API for 3-Floyd attention.
+- Added additional configuration options in `PivotalAttentionBlock`.
+
+The full changelog is in [CHANGELOG.md](CHANGELOG.md).
+
+## Citation
+
+If you use this code in your research, please cite the paper:
+
+```bibtex
+@inproceedings{TODO,
+  title     = {TODO},
+  author    = {TODO},
+  booktitle = {International Conference on Learning Representations (ICLR)},
+  year      = {TODO},
+  url       = {TODO}
+}
+```
+
+(Alternatively, see [CITATION.cff](CITATION.cff).)
+
+---
 
 ## License
-MIT
+
+This project is licensed under the **Apache License 2.0**. See [LICENSE](LICENSE).

floydnet-0.1.1/README.md

@@ -0,0 +1,159 @@
+# FloydNet
+
+Official implementation of an ICLR paper (TODO: add paper title, authors, and links/arXiv).
+
+![Figure Pivotal Attention Mechanism for 2-Floyd/3-Floyd.](misc/pivotalattn2&3.png)
+
+This repository serves two audiences:
+
+- **Engineering users**: reusable PyTorch components (functional attention APIs and Transformer-style blocks) under `src/`.
+- **Research users**: scripts/configs to reproduce paper experiments under `example/`.
+
+---
+
+## Introduction
+
+FloydNet is the official PyTorch implementation accompanying an ICLR paper (TODO).  
+The repository provides:
+
+1. **Reusable components**: a drop-in attention/Transformer-block interface intended for integration into existing projects.
+2. **Reproduction code**: end-to-end training/evaluation pipelines to reproduce the benchmarks reported in the paper.
+
+For algorithmic details, hyperparameter choices, and analysis, please refer to the paper (TODO: link).
+
+---
+
+## Repository Structure
+
+- `src/floydnet/`  
+  **Library code for reuse**  
+  Contains the functional attention API and module/block implementations.
+
+- `example/`  
+  **Experiment reproduction code**  
+  Includes benchmark-specific scripts, configs, and data preparation utilities.
+
+---
+
+## Using the Attention / Transformer Block
+
+This section targets **engineering users** who want to import FloydNet as a dependency.
+
+### Installation
+
+#### Option A: Install from PyPI
+```bash
+pip install floydnet
+```
+
+#### Option B: Install from source
+```bash
+git clone git@github.com:ocx-lab/FloydNet.git
+cd FloydNet
+pip install -e .
+```
+
+> Requirements: Python `>= 3.9`, PyTorch `>= 2.1` (see `pyproject.toml`).
+
+### Public API
+
+FloydNet re-exports the public API from `src/floydnet/__init__.py`, so you can import from the top-level package:
+
+- **Functional API**:
+  - `pivotal_attention` (see `src/floydnet/functional.py`)
+- **Module / block API**:
+  - `PivotalAttentionBlock` (see `src/floydnet/transformer.py`)
+
+```python
+from floydnet import pivotal_attention, PivotalAttentionBlock
+```
+
+### Minimal usage example
+
+```python
+import torch
+from floydnet import pivotal_attention, PivotalAttentionBlock
+
+# -------------------------
+# Module API (Transformer-style block)
+# Input is a 2D grid: (B, N, N, C)
+# -------------------------
+B, N, C = 2, 16, 64
+x = torch.randn(B, N, N, C)
+
+m = PivotalAttentionBlock(embed_dim=C, num_heads=8, dropout=0.0)
+out = m(x)  # (B, N, N, C)
+print(out.shape)
+
+# -------------------------
+# Functional API
+# All inputs are 5D: (B, H, N, N, D)
+# -------------------------
+B, H, N, D = 2, 8, 16, 64
+q_ik = torch.randn(B, H, N, N, D)
+k_ij = torch.randn(B, H, N, N, D)
+k_jk = torch.randn(B, H, N, N, D)
+v_ij = torch.randn(B, H, N, N, D)
+v_jk = torch.randn(B, H, N, N, D)
+
+y = pivotal_attention(q_ik, k_ij, k_jk, v_ij, v_jk)  # (B, H, N, N, D)
+print(y.shape)
+```
+
+---
+
+## Reproducing Paper Results
+
+This section targets **research users** who want to reproduce the experiments in the paper.
+
+See `example/README.md` For detailed description.
+
+### Environment setup
+
+We recommend using `uv` to create an isolated environment for the reproduction code under `example/`.
+
+```bash
+cd /path/to/FloydNet
+
+# 1) Create a uv virtual environment with Python 3.12
+uv venv --python 3.12
+
+# 2) Activate it
+source .venv/bin/activate
+
+# 3) Install extra dependencies for reproducing paper experiments
+uv pip install -r example/requirements.txt
+
+# 4) Install FloydNet (editable) for local development / imports
+uv pip install -e .
+```
+
+## Changelog (latest)
+
+- Full release with training and evaluation scripts for Graph Count, BREC, and TSP.
+- Added `pivotal_attention3` functional API for 3-Floyd attention.
+- Added additional configuration options in `PivotalAttentionBlock`.
+
+The full changelog is in [CHANGELOG.md](CHANGELOG.md).
+
+## Citation
+
+If you use this code in your research, please cite the paper:
+
+```bibtex
+@inproceedings{TODO,
+  title     = {TODO},
+  author    = {TODO},
+  booktitle = {International Conference on Learning Representations (ICLR)},
+  year      = {TODO},
+  url       = {TODO}
+}
+```
+
+(Alternatively, see [CITATION.cff](CITATION.cff).)
+
+---
+
+## License
+
+This project is licensed under the **Apache License 2.0**. See [LICENSE](LICENSE).

floydnet-0.1.1/example/README.md

@@ -0,0 +1,129 @@
+### Benchmarks
+
+The paper reports results on **three benchmarks**:
+
+- Graph Count
+- BREC
+- TSP
+
+### Graph Count
+
+The Graph Count benchmark and dataset construction follow:
+https://github.com/subgraph23/homomorphism-expressivity
+
+```bash
+source .venv/bin/activate
+cd example
+python -m data.count.process_data
+./count/run.sh
+```
+
+### BREC
+
+The BREC benchmark and dataset construction follow:
+https://github.com/GraphPKU/BREC
+
+```bash
+source .venv/bin/activate
+cd example/data/BREC/raw
+unzip BREC_data_all.zip
+
+# Back to the example folder
+cd ../../..
+
+# Reproduce FloydNet results
+python -m BREC.test_BREC
+
+# Reproduce 3-Floyd results
+python -m BREC.test_BREC --floyd_level 3
+```
+
+### TSP
+
+Reproducing TSP at full scale is computationally heavy and involves large datasets. For convenience, we provide:
+
+- A small demo dataset on Hugging Face:  
+  https://huggingface.co/datasets/ocxlabs/FloydNet_TSP_demo
+- Pretrained checkpoints for:
+  - **Metric TSP** (Euclidean TSP, `euc`): https://huggingface.co/ocxlabs/FloydNet_TSP_euc
+  - **Non-metric TSP** (Explicit TSP, `exp`): https://huggingface.co/ocxlabs/FloydNet_TSP_exp
+
+This section describes **inference and evaluation** using the demo data and checkpoints.
+
+#### Prepare demo data
+
+Download the demo dataset as a `.zip`, unzip it, and place the extracted folder under `example/data/`.
+
+#### Inference
+
+Run inference in `--test_mode` using `torchrun` (the command below assumes **single-node, 8 GPUs**).  
+Set `--subset` and make sure `--load_checkpoint` matches the subset.
+
+```bash
+source .venv/bin/activate
+cd example
+
+torchrun \
+  --nproc_per_node=8 \
+  -m TSP.run \
+  --subset exp \
+  --output_dir ./outputs/TSP_exp \
+  --load_checkpoint path/to/TSP_exp/epoch_01000.pt \
+  --test_mode \
+  --split_factor 1 \
+  --sample_count_per_case 10
+```
+
+#### Evaluation
+
+After inference finishes, aggregate results with:
+
+```bash
+source .venv/bin/activate
+cd example
+
+python -m TSP.report ./outputs/TSP_exp
+```
+
+This saves CSV summaries (downsampled to 1 / 5 / 10 samples per instance) into the same `output_dir`.
+
+#### Data generation
+
+If you want to generate additional data (beyond the demo set) and train from scratch, prepare the raw `.npy` files as follows.
+
+##### Metric TSP (Euclidean TSP, `euc`)
+
+1. Randomly sample **N integer points** in 2D, $p_i = (x_i, y_i)$, and ensure **pairwise Euclidean distances ≤ 200**.
+2. Solve the instance with a classic TSP solver (e.g., [Concorde](https://www.math.uwaterloo.ca/tsp/concorde.html)).
+3. Reorder the points so that $p_0 \rightarrow p_1 \rightarrow ... \rightarrow p_{N-1}$ is the optimal tour.
+4. Sample **T** instances for each **N**, stack them into a NumPy array of shape **`[T, N, 2]`** with dtype **`int8`**, and save grouped-by-N arrays as:
+   - `data/TSP/euc/non-uni/raw/{N:03d}.npy`
+
+##### Non-metric TSP (Explicit TSP, `exp`)
+
+1. Randomly sample an **N×N symmetric distance matrix** with **maximum value ≤ 200**.
+2. Solve with a classic TSP solver (e.g., [Concorde](https://www.math.uwaterloo.ca/tsp/concorde.html)).
+3. Reorder rows/columns so that $0 \rightarrow 1 \rightarrow ... \rightarrow N-1$ is the optimal tour.
+4. Sample **T** instances for each **N**, stack them into a NumPy array of shape **`[T, N, N]`** with dtype **`int8`**, and save grouped-by-N arrays as:
+   - `data/TSP/exp/non-uni/raw/{N:03d}.npy`
+
+#### Training from scratch
+
+Recommended (paper-matching) training setup: 8 nodes × 8 GPUs = 64 GPUs total.
+
+```bash
+source .venv/bin/activate
+cd example
+torchrun \
+  --master_addr <MASTER_ADDR> \
+  --master_port <MASTER_PORT> \
+  --nnodes 8 \
+  --node_rank <NODE_RANK> \
+  --nproc_per_node 8 \
+  -m TSP.run \
+  --subset exp \
+  --output_dir ./outputs/TSP_exp \
+  --wandb_name TSP_exp
+```
+
+---

{floydnet-0.1.0 → floydnet-0.1.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "floydnet"
-version = "0.1.0"
+version = "0.1.1"
 description = "Floyd Multi-Head Attention: a drop-in variant of PyTorch MHA with module and function APIs"
 readme = "README.md"
 requires-python = ">=3.9"
@@ -54,7 +54,7 @@ include = [
 ]
 
 [tool.hatch.build.targets.wheel]
-packages = ["src/floyd_net"]
+packages = ["src/floydnet"]
 
 [tool.ruff]
 line-length = 100

floydnet-0.1.0/CHANGELOG.md

@@ -1,7 +0,0 @@
-# Changelog
-
-All notable changes to this project will be documented in this file.
-
-## [0.1.0] - 2025-10-21
-- Initial public skeleton with module + functional attention
-- CI, tests, examples, paper scaffolding

floydnet-0.1.0/README.md

@@ -1,46 +0,0 @@
-# floyd-net
-
-Floyd Multi-Head Attention (F-MHA) is a drop-in variant of PyTorch's attention stack. It provides:
-
-- Module API: `FloydMultiheadAttention` mirroring `torch.nn.MultiheadAttention`
-- Functional API: `floyd_scaled_dot_product_attention` mirroring `torch.nn.functional.scaled_dot_product_attention`
-
-Install and manage with `uv` for a modern Python workflow.
-
-## Quick start
-
-```bash
-# Install with uv (recommended)
-uv venv --python 3.10
-source .venv/bin/activate
-uv pip install -e .[dev]
-```
-
-```python
-import torch
-from floyd_net import FloydMultiheadAttention
-
-m = FloydMultiheadAttention(embed_dim=64, num_heads=8, batch_first=True)
-x = torch.randn(2, 16, 64)
-out, attn = m(x, x, x)
-print(out.shape)
-```
-
-### Functional API
-```python
-import torch
-import torch.nn.functional as F
-from floyd_net import floyd_scaled_dot_product_attention
-
-q = torch.randn(2, 8, 16, 64)  # (B, H, L, D)
-k = torch.randn(2, 8, 16, 64)
-v = torch.randn(2, 8, 16, 64)
-out = floyd_scaled_dot_product_attention(q, k, v)
-print(out.shape)
-```
-
-## Paper reproductions
-See `paper/` for dataset preparation, configs, and experiment templates to reproduce the results in the paper.
-
-## License
-MIT

floydnet-0.1.0/paper/README.md

@@ -1,11 +0,0 @@
-# Paper reproductions
-
-This folder contains materials to reproduce results from the paper.
-
-Structure:
-- `datasets/`: dataset preparation scripts or links
-- `configs/`: YAML/TOML experiment configs
-- `experiments/`: runnable training/eval scripts referencing configs
-- `notebooks/`: exploratory notebooks (optional)
-
-Use `uv` to create an environment, then run scripts inside `experiments/`.