floydnet 0.1.0__tar.gz → 0.1.2__tar.gz

{floydnet-0.1.0 → floydnet-0.1.2}/.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.2/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.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: floydnet
-Version: 0.1.0
+Version: 0.1.2
 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,160 @@ Requires-Dist: pytest>=7.4; extra == 'dev'
 Requires-Dist: ruff>=0.5; extra == 'dev'
 Description-Content-Type: text/markdown
 
-# floyd-net
+# FloydNet
+[![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
+[![Python](https://img.shields.io/badge/Python-3.9%2B-blue)](https://www.python.org/)
+[![PyTorch](https://img.shields.io/badge/PyTorch-2.1%2B-orange)](https://pytorch.org/)
 
-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:
+- **Engineering users**: Reusable PyTorch components (functional attention APIs and Transformer-style blocks) under `src/`.
+- **Research users**: Scripts/configs to reproduce paper experiments (TSP, Graph Isomorphism, BREC) under `example/`.
 
-## Quick start
+---
 
+## 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.
+
+---
+
+### 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)
+```
+
+---
+
+## 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}
+}
 ```
 
-## Paper reproductions
-See `paper/` for dataset preparation, configs, and experiment templates to reproduce the results in the paper.
+(Alternatively, see [CITATION.cff](CITATION.cff).)
+
+---
 
 ## License
-MIT
+
+This project is licensed under the **Apache License 2.0**. See [LICENSE](LICENSE).

floydnet-0.1.2/README.md

@@ -0,0 +1,157 @@
+# FloydNet
+[![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
+[![Python](https://img.shields.io/badge/Python-3.9%2B-blue)](https://www.python.org/)
+[![PyTorch](https://img.shields.io/badge/PyTorch-2.1%2B-orange)](https://pytorch.org/)
+
+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 (TSP, Graph Isomorphism, BREC) 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.
+
+---
+
+### 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.2/example/README.md

@@ -0,0 +1,137 @@
+### Benchmarks
+
+The paper reports results on **three benchmarks**:
+
+- Graph Count
+- BREC
+- TSP
+
+## 🚀 Key Results
+
+| Domain | Benchmark | Result |
+| :--- | :--- | :--- |
+| **Algorithmic** | CLRS-30 | **96.64%** (SOTA), effectively solving graph & string algorithms. |
+| **Optimization** | Non-Metric TSP | **88.3%** optimality on unseen sizes ($N=200$), vs 1.3% for Linkern heuristic. |
+| **Expressiveness** | Substructure Count | Near-zero error (MAE **0.001**) on complex substructure counting. |
+
+### 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.2}/pyproject.toml

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

floydnet-0.1.2/src/floydnet/__init__.py

@@ -0,0 +1,8 @@
+from .functional import pivotal_attention, pivotal_attention3
+from .transformer import PivotalAttentionBlock
+
+__all__ = [
+    "pivotal_attention",
+    "pivotal_attention3",
+    "PivotalAttentionBlock",
+]

floydnet-0.1.2/src/floydnet/functional.py

@@ -0,0 +1,150 @@
+# Copyright 2025 Beijing Academy of Artificial Intelligence (BAAI)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from __future__ import annotations
+
+from typing import Optional
+import math
+
+import torch
+import torch.nn.functional as F
+
+def pivotal_attention(
+    q_ik: torch.Tensor,
+    k_ij: torch.Tensor,
+    k_jk: torch.Tensor,
+    v_ij: torch.Tensor,
+    v_jk: torch.Tensor,
+    attn_mask: Optional[torch.Tensor] = None,
+    dropout: float = 0.0,
+    scale: Optional[float] = None,
+    inf: float = 1e9,
+) -> torch.Tensor:
+    """Pivotal attention as described in "FLOYDNET: A LEARNING PARADIGM FOR GLOBAL RELATIONAL REASONING".
+
+    Shapes:
+        q_ik: (B, H, L_i, L_k, D)
+        k_ij: (B, H, L_i, L_j, D)
+        k_jk: (B, H, L_j, L_k, D)
+        v_ij: (B, H, L_i, L_j, D)
+        v_jk: (B, H, L_j, L_k, D)
+        attn_mask (optional): broadcastable to (B, H, L_i, L_k, L_j)
+
+    Args:
+        attn_mask: Additive mask (float) or boolean mask. If boolean, masked positions are set to -inf.
+        dropout: Dropout probability applied to attention weights (only effective if > 0).
+        scale: Optional custom scaling factor. If None, defaults to 1/sqrt(2*D).
+        inf: Value to use for -infinity in masks.
+
+    Returns:
+        Tensor of shape (B, H, L_i, L_k, D)
+    """
+    assert all([t.dim() == 5 for t in [q_ik, k_ij, k_jk, v_ij, v_jk]]), "All inputs must be 5D tensors"
+    B, H, L_i, L_k, D = q_ik.shape
+    L_j = k_ij.shape[3]
+    assert k_ij.shape == v_ij.shape == (B, H, L_i, L_j, D), "k_ij and v_ij must have shape (B, H, L_i, L_j, D)"
+    assert k_jk.shape == v_jk.shape == (B, H, L_j, L_k, D), "k_jk and v_jk must have shape (B, H, L_j, L_k, D)"
+
+    if scale is None:
+        scale = 1.0 / math.sqrt(2.0 * D)
+    q_ik = q_ik * scale
+
+    # Compute attention scores over the pivot dimension j: (B, H, L_i, L_k, L_j)
+    attn_scores = torch.einsum("bhikd,bhijd->bhikj", q_ik, k_ij) \
+                + torch.einsum("bhikd,bhjkd->bhikj", q_ik, k_jk)
+
+    if attn_mask is not None:
+        if attn_mask.dtype == torch.bool:
+            attn_scores = attn_scores.masked_fill(attn_mask, -inf)
+        else:
+            attn_scores = attn_scores + attn_mask
+
+    attn_weights = torch.softmax(attn_scores, dim=-1)
+
+    if dropout > 0.0:
+        attn_weights = F.dropout(attn_weights, p=dropout)
+
+    y = torch.einsum("bhikj,bhijd->bhikd", attn_weights, v_ij) \
+      + torch.einsum("bhikj,bhjkd->bhikd", attn_weights, v_jk)
+
+    return y
+
+def pivotal_attention3(
+    q_ijk: torch.Tensor,
+    k_pjk: torch.Tensor,
+    k_ipk: torch.Tensor,
+    k_ijp: torch.Tensor,
+    v_pjk: torch.Tensor,
+    v_ipk: torch.Tensor,
+    v_ijp: torch.Tensor,
+    attn_mask: Optional[torch.Tensor] = None,
+    dropout: float = 0.0,
+    scale: Optional[float] = None,
+    inf: float = 1e9,
+) -> torch.Tensor:
+    """3-Pivotal attention as described in "FLOYDNET: A LEARNING PARADIGM FOR GLOBAL RELATIONAL REASONING".
+
+    Shapes:
+        q_ijk: (B, H, L_i, L_j, L_k, D)
+        k_pjk: (B, H, L_p, L_j, L_k, D)
+        k_ipk: (B, H, L_i, L_p, L_k, D)
+        k_ijp: (B, H, L_i, L_j, L_p, D)
+        v_pjk: (B, H, L_p, L_j, L_k, D)
+        v_ipk: (B, H, L_i, L_p, L_k, D)
+        v_ijp: (B, H, L_i, L_j, L_p, D)
+        attn_mask (optional): broadcastable to (B, H, L_i, L_j, L_k, L_p)
+
+    Args:
+        attn_mask: Additive mask (float) or boolean mask. If boolean, masked positions are set to -inf.
+        dropout: Dropout probability applied to attention weights (only effective if > 0).
+        scale: Optional custom scaling factor. If None, defaults to 1/sqrt(3*D).
+        inf: Value to use for -infinity in masks.
+
+    Returns:
+        Tensor of shape (B, H, L_i, l_j, L_k, D)
+    """
+    assert all([t.dim() == 6 for t in [q_ijk, k_pjk, k_ipk, k_ijp, v_pjk, v_ipk, v_ijp]]), "All inputs must be 6D tensors"
+    B, H, L_i, L_j, L_k, D = q_ijk.shape
+    L_p = k_pjk.shape[2]
+    assert k_pjk.shape == v_pjk.shape == (B, H, L_p, L_j, L_k, D), "k_pjk and v_pjk must have shape (B, H, L_p, L_j, L_k, D)"
+    assert k_ipk.shape == v_ipk.shape == (B, H, L_i, L_p, L_k, D), "k_ipk and v_ipk must have shape (B, H, L_i, L_p, L_k, D)"
+    assert k_ijp.shape == v_ijp.shape == (B, H, L_i, L_j, L_p, D), "k_ijp and v_ijp must have shape (B, H, L_i, L_j, L_p, D)"
+    
+    if scale is None:
+        scale = 1.0 / math.sqrt(3.0 * D)
+    q_ijk = q_ijk * scale
+
+    # Compute attention scores over the pivot dimension j: (B, H, L_i, L_j, L_k, L_p)
+    attn_scores = torch.einsum("bhijkd,bhpjkd->bhijkp", q_ijk, k_pjk) \
+                + torch.einsum("bhijkd,bhipkd->bhijkp", q_ijk, k_ipk) \
+                + torch.einsum("bhijkd,bhijpd->bhijkp", q_ijk, k_ijp)
+
+    if attn_mask is not None:
+        if attn_mask.dtype == torch.bool:
+            attn_scores = attn_scores.masked_fill(attn_mask, -inf)
+        else:
+            attn_scores = attn_scores + attn_mask
+
+    attn_weights = torch.softmax(attn_scores, dim=-1)
+
+    if dropout > 0.0:
+        attn_weights = F.dropout(attn_weights, p=dropout)
+
+    y = torch.einsum("bhijkp,bhpjkd->bhijkd", attn_weights, v_pjk) \
+      + torch.einsum("bhijkp,bhipkd->bhijkd", attn_weights, v_ipk) \
+      + torch.einsum("bhijkp,bhijpd->bhijkd", attn_weights, v_ijp)
+
+    return y
+

floydnet-0.1.2/src/floydnet/transformer.py

@@ -0,0 +1,219 @@
+# Copyright 2025 Beijing Academy of Artificial Intelligence (BAAI)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import copy
+from typing import Optional, Tuple, Union, Callable
+
+import torch
+from torch import nn
+from .functional import pivotal_attention
+
+
+class Affine(nn.Module):
+    def __init__(self, c):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones((c, )))
+        self.bias = nn.Parameter(torch.zeros((c, )))
+
+    def forward(self, x: torch.Tensor):
+        return x * self.weight + self.bias
+
+
+def create_norm(norm_fn: Union[str, Callable], embed_dim: int, eps: float = 1e-5, **kwargs) -> nn.Module:
+    """Create a normalization module from a name or nn.Module.
+
+    Args:
+        norm_fn: Name or an nn.Module instance/class.
+        embed_dim: Embedding dimension (features) used to construct the norm.
+        eps: Numerical epsilon passed to the normalization layer if applicable.
+        **kwargs: Extra keyword arguments forwarded to the normalization layer.
+
+    Returns:
+        An nn.Module normalization instance.
+    """
+    if isinstance(norm_fn, str):
+        if norm_fn.lower() in ["layernorm", "ln"]:
+            return nn.LayerNorm(embed_dim, eps=eps, **kwargs)
+        elif norm_fn.lower() in ["batchnorm", "bn"]:
+            return nn.BatchNorm1d(embed_dim, eps=eps, **kwargs)
+        elif norm_fn.lower() in ["rmsnorm", "rms"]:
+            return nn.RMSNorm(embed_dim, eps=eps, **kwargs)
+        elif norm_fn.lower() in ["affine"]:
+            return Affine(embed_dim)
+        elif norm_fn.lower() in ["none", "identity"]:
+            return nn.Identity()
+        else:
+            raise ValueError(f"Unsupported norm_fn string: {norm_fn}")
+    elif callable(norm_fn):
+        if isinstance(norm_fn, nn.Module):
+            # deepcopy to avoid shared parameters
+            return copy.deepcopy(norm_fn)
+        elif isinstance(norm_fn, type) and issubclass(norm_fn, nn.Module):
+            return norm_fn(embed_dim, eps=eps, **kwargs)
+        else:
+            raise TypeError("norm_fn callable must be an nn.Module or nn.Module class")
+    else:
+        raise TypeError("norm_fn must be a string or callable")
+
+        
+def create_activation(activation_fn: Union[str, Callable]) -> nn.Module:
+    """Create an activation module from a name or nn.Module.
+
+    Args:
+        activation_fn: Name or an nn.Module instance/class.
+
+    Returns:
+        An nn.Module activation instance.
+    """
+    if isinstance(activation_fn, str):
+        if activation_fn.lower() == "relu":
+            return nn.ReLU()
+        elif activation_fn.lower() == "gelu":
+            return nn.GELU()
+        elif activation_fn.lower() == "silu":
+            return nn.SiLU()
+        else:
+            raise ValueError(f"Unsupported activation_fn string: {activation_fn}")
+    elif callable(activation_fn):
+        if isinstance(activation_fn, nn.Module):
+            return activation_fn
+        elif isinstance(activation_fn, type) and issubclass(activation_fn, nn.Module):
+            return activation_fn()
+        else:
+            raise TypeError("activation_fn callable must be an nn.Module or nn.Module class")
+    else:
+        raise TypeError("activation_fn must be a string or callable")
+
+
+class PivotalAttentionBlock(nn.Module):
+    """Transformer-style block that applies pivotal attention followed by an FFN.
+
+    Args:
+        embed_dim: Input/hidden embedding dimension (D).
+        num_heads: Number of attention heads (D must be divisible by num_heads).
+        dropout: Dropout probability for attention output and FFN output.
+        bias: Whether to include bias terms in linear layers.
+        ffn_expansion_ratio: Expansion ratio for the FFN hidden size.
+        norm_position: "pre" or "post" layer normalization placement.
+        activation_fn: Activation name/module used in the FFN.
+        norm_fn: Normalization name/module used in the block.
+    """
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        dropout: float = 0.0,
+        bias: bool = False,
+        ffn_expansion_ratio: int = 4,
+        norm_position: str = "pre",
+        activation_fn: Union[str, Callable] = "relu",
+        norm_fn: Union[str, Callable] = "layernorm",
+        enable_symmetric_mix: bool = True,
+        enable_ffn: bool = True,
+    ) -> None:
+        super().__init__()
+        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        self.dropout = dropout
+        self.norm_position = norm_position.lower()
+        self.enable_ffn = enable_ffn
+        assert self.norm_position in ["pre", "post"], "norm_position must be 'pre' or 'post'"
+
+        self.enable_symmetric_mix = enable_symmetric_mix
+        if enable_symmetric_mix:
+            self.c_mix = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+        self.c_qkv = nn.Linear(embed_dim, embed_dim * 5, bias=bias)
+        self.c_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.dropout_fn = nn.Dropout(dropout)
+        self.norm1 = create_norm(norm_fn, embed_dim)
+        if self.enable_ffn:
+            self.activation_fn = create_activation(activation_fn)
+            self.norm2 = create_norm(norm_fn, embed_dim)
+            self.ffn = nn.Sequential(
+                nn.Linear(embed_dim, ffn_expansion_ratio * embed_dim, bias=bias),
+                self.activation_fn,
+                nn.Linear(ffn_expansion_ratio * embed_dim, embed_dim, bias=bias),
+                nn.Dropout(dropout),
+            )
+            self.ffn_scale = nn.Parameter(torch.tensor(1.0, requires_grad=True))
+
+        self._reset_parameters()
+
+    def _reset_parameters(self) -> None:
+        """Initialize parameters using Xavier for projections and zeros for output heads."""
+        if self.enable_symmetric_mix:
+            nn.init.zeros_(self.c_mix.weight)
+        nn.init.xavier_uniform_(self.c_qkv.weight)
+        nn.init.zeros_(self.c_proj.weight)
+        if self.enable_ffn:
+            nn.init.xavier_uniform_(self.ffn[0].weight)
+            nn.init.zeros_(self.ffn[2].weight)
+        if self.c_qkv.bias is not None:
+            if self.enable_symmetric_mix:
+                nn.init.zeros_(self.c_mix.bias)
+            nn.init.zeros_(self.c_qkv.bias)
+            nn.init.zeros_(self.c_proj.bias)
+            if self.enable_ffn:
+                nn.init.zeros_(self.ffn[0].bias)
+                nn.init.zeros_(self.ffn[2].bias)
+
+    def attn(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor]) -> torch.Tensor:
+        """Apply pivotal attention over a (L x L) grid.
+
+        Args:
+            x: Input tensor of shape (B, L, L, D).
+            attn_mask: Optional mask broadcastable to (B, H, L, L, L).
+
+        Returns:
+            Tensor of shape (B, L, L, D) after attention projection and dropout.
+        """
+        B, L, _, D = x.shape
+        # [B, L, L, 5*D] -> 5 x [B, H, L, L, d]
+        qkv = torch.chunk(self.c_qkv(x), 5, dim=-1)
+        q_ik, k_ij, k_jk, v_ij, v_jk = map(
+            lambda t: t.view(B, L, L, self.num_heads, self.head_dim).permute(0, 3, 1, 2, 4),
+            qkv, 
+        )
+
+        # [B, H, L, L, d]
+        y = pivotal_attention(
+            q_ik, k_ij, k_jk, v_ij, v_jk,
+            attn_mask=attn_mask,
+            dropout=self.dropout if self.training else 0.0,
+        )
+        y = y.permute(0, 2, 3, 1, 4).contiguous().view(B, L, L, D)
+        y = self.c_proj(y)
+        y = self.dropout_fn(y)
+        return y
+
+    def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        if self.enable_symmetric_mix:
+            xT = self.c_mix(x.transpose(1, 2))
+        else:
+            xT = 0
+        if self.norm_position == "pre":
+            x = x + self.attn(self.norm1(x + xT), attn_mask)
+            if self.enable_ffn:
+                x = x + self.ffn(self.norm2(x)) * self.ffn_scale
+        else:
+            x = self.norm1(x + self.attn(x + xT, attn_mask))
+            if self.enable_ffn:
+                x = self.norm2(x + self.ffn(x)) * self.ffn_scale
+
+        return x

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/`.