merge to main for 0.1.0

.git-blame-ignore-revs

@@ -0,0 +1,5 @@
+# Commits to ignore in git blame (formatting-only changes).
+# GitHub picks this file up automatically; locally, enable it with:
+#   git config blame.ignoreRevsFile .git-blame-ignore-revs
+
+eb65ffb809108b784e2051767f0c0f3f9e672fb3

.pre-commit-config.yaml

@@ -1,19 +1,20 @@
+exclude: (^examples\/)
 repos:
-- repo: https://github.com/psf/black
-  rev: 25.1.0
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  rev: v0.15.20
   hooks:
-  - id: black
-    language_version: python3
-    args: ['--line-length=125',
-           '--skip-string-normalization']
+    - id: ruff
+      args: [--fix]
+    - id: ruff-format
 
 - repo: https://github.com/tox-dev/pyproject-fmt
-  rev: v2.5.1
+  rev: v2.25.1
   hooks:
     - id: pyproject-fmt
+      args: ["--max-supported-python", "3.13"]
 
 - repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v5.0.0
+  rev: v6.0.0
   hooks:
   - id: check-added-large-files
   - id: check-case-conflict
@@ -27,29 +28,26 @@ repos:
   - id: requirements-txt-fixer
   - id: trailing-whitespace
 
-- repo: https://github.com/PyCQA/isort
-  rev: 6.0.1
-  hooks:
-  - id: isort
-    args: ["--profile=black"] # <-- this one
-
 - repo: https://github.com/asottile/pyupgrade
-  rev: v3.19.1
+  rev: v3.21.2
   hooks:
   - id: pyupgrade
-    args: ["--py36-plus"]
+    args: ["--py310-plus"]
 
 - repo: https://github.com/pycqa/flake8
-  rev: 7.1.2
+  rev: 7.3.0
   hooks:
   - id: flake8
-    exclude: docs/conf.py
+    exclude: docs/source/conf.py
     additional_dependencies: [flake8-bugbear, flake8-print]
     args: ['--max-line-length=125',  # github viewer width
-           '--extend-ignore=E203']  # E203 is not PEP8 compliant
+           '--extend-ignore=E203,T201,F401',
+           # E203 is not PEP8 compliant
+           # F401 included in ruff (behaves slightly differently for noqa flags)
+    ]
 
 - repo: https://github.com/mgedmin/check-manifest
-  rev: "0.50"
+  rev: "0.51"
   hooks:
   - id: check-manifest
     stages: [manual]

.readthedocs.yaml

@@ -12,12 +12,11 @@ build:
 
 # Build documentation in the "docs/" directory with Sphinx
 sphinx:
-   configuration: docs/conf.py
+   configuration: docs/source/conf.py
 
 # Optionally, but recommended,
 # declare the Python requirements required to build your documentation
 # See https://docs.readthedocs.io/en/stable/guides/reproducible-builds.html
-# python:
-#    install:
-#    - requirements: docs/requirements.txt
-
+python:
+  install:
+    - requirements: docs/requirements.txt

README.md

@@ -1,61 +1,118 @@
-![alt text](docs/source/_static/pquant_white_font.png)
+<p align="center">
+  <img src="docs/source/_static/pquant_white_font.png" alt="PQuantML logo" width="400">
+</p>
 
-## Prune and Quantize ML models
-PQuant is a library for training compressed machine learning models, developed at CERN as part of the [Next Generation Triggers](https://nextgentriggers.web.cern.ch/t13/) project.
+# PQuantML
 
-Installation via pip: ```pip install pquant-ml```. 
+**PQuantML** is an end-to-end library for training compressed machine learning models, developed at CERN as part of the [Next Generation Triggers](https://nextgentriggers.web.cern.ch/t13/) project.
 
-With TensorFlow ```pip install pquant-ml[tensorflow]```. 
+It supports:
 
-With PyTorch ```pip install pquant-ml[torch]```.
+- Joint pruning + quantization
+- Layer-wise precision configuration
+- Flexible training pipelines
+- PyTorch and TensorFlow backends
+- Knowledge distillation
+- HGQ library integration
+- Integration with hardware-friendly toolchains (e.g., [hls4ml](https://fastmachinelearning.org/hls4ml/))
 
-PQuant replaces the layers and activations it finds with a Compressed (in the case of layers) or Quantized (in the case of activations) variant. These automatically handle the quantization of the weights, biases and activations, and the pruning of the weights. 
-Both PyTorch and TensorFlow models are supported. 
+PQuantML enables efficient deployment of compact neural networks on resource-constrained hardware such as FPGAs and embedded accelerators.
 
-### Layers that can be compressed
+<p align="center">
+  <img src="docs/source/_static/overview_pquant_updated.png" alt="PQuantML overview" width="1000">
+</p>
 
-* **PQConv*D**: Convolutional layers
-* **PQAvgPool*D**: Average pooling layers
-* **PQBatchNorm*D**: BatchNorm layers
-* **PQDense**: Linear layer
-* **PQActivation**: Activation layers (ReLU, Tanh)
+## Installation
 
-The various pruning methods have different training steps, such as a pre-training step and fine-tuning step. PQuant provides a training function, where the user provides the functions to train and validate an epoch, and PQuant handles the training while triggering the different training steps.
+Install the base package via pip:
 
+```bash
+pip install pquant-ml
+```
 
-![alt text](docs/source/_static/overview_pquant.png)
+Install with a specific backend:
 
+```bash
+pip install "pquant-ml[tensorflow]"   # TensorFlow backend
+pip install "pquant-ml[torch]"        # PyTorch backend
+```
 
+## Supported layers
 
-### Example
-Example notebook can be found [here](https://github.com/nroope/PQuant/tree/main/examples). It handles the
-  1. Creation of a torch model and data loaders.
-  2. Creation of the training and validation functions.
-  3. Loading a default pruning configuration of a pruning method.
-  4. Using the configuration, the model, and the training and validation functions, call the training function of PQuant to train and compress the model.
-  5. Creating a custom quantization and pruning configuration for a given model (disable pruning for some layers, different quantization bitwidths for different layers).
-  6. Direct layers usage and layers replacement approaches.
-  7. Usage of fine-tuning platform.
+| Layer | Description |
+| --- | --- |
+| `PQConv*D` | Convolutional layers |
+| `PQAvgPool*D` | Average pooling layers |
+| `PQBatchNorm*D` | Batch normalization layers |
+| `PQDense` | Linear (fully connected) layer |
+| `PQActivation` | Activation layers: ReLU, Tanh, Leaky ReLU, GELU, Hard Tanh, or a user-provided activation function (Torch only) |
+| `MultiHeadAttention` | Multi-head attention layer |
+| `LayerNorm` | Layer normalization layer (currently Torch only) |
 
-### Pruning methods
-A description of the pruning methods and their hyperparameters can be found [here](docs/pruning_methods.md).
+## Training
 
-### Quantization parameters
-A description of the quantization parameters can be found [here](docs/quantization_parameters.md).
+Different pruning methods involve different training stages, such as pre-training and fine-tuning. PQuantML provides a generic training function: you supply your own training and validation functions along with the number of epochs, and PQuant handles the training loop while automatically triggering the appropriate stages for the chosen pruning method.
+<p align="center">
+  <img src="docs/source/_static/pruning_methods_overview.png" alt="Pruning methods overview." width="1000">
+</p>
 
+## Quantization
 
-For detailed documentation check this page: [PQuantML documentation](https://pquantml.readthedocs.io/en/latest/)
+PQuantML supports two quantization modes, each with several granularity options.
 
+**Fixed-point quantization** (for weights):
+- per-weight
+- per-channel
+- per-tensor
 
-### Authors
- - Roope Niemi (CERN)
- - Anastasiia Petrovych (CERN)
- - Arghya Das (Purdue University)
- - Enrico Lupi (CERN)
- - Chang Sun (Caltech)
- - Dimitrios Danopoulos (CERN)
- - Marlon Joshua Helbing
- - Mia Liu (Purdue University)
- - Michael Kagan (SLAC National Accelerator Laboratory)
- - Vladimir Loncar (CERN)
- - Maurizio Pierini (CERN)
+**HGQ (High Granularity Quantization):**
+- per-weight
+- per-tensor
+
+## Example
+
+Example notebooks are available in the [`examples/`](https://github.com/cern-nextgen/PQuantML/tree/main/examples) directory. It shows how to:
+
+1. Create a Torch model and data loaders.
+2. Create the training and validation functions.
+3. Load a default configuration for a pruning method.
+4. Train and compress the model by passing the configuration, model, and training/validation functions to PQuant's training function.
+5. Build a custom quantization and pruning configuration for a given model (e.g. disabling pruning for some layers, or using different quantization bit-widths per layer).
+6. Use the direct-layer and layer-replacement approaches.
+7. Use the HPO platform.
+
+## Documentation
+
+Full documentation is available at [pquantml.readthedocs.io](https://pquantml.readthedocs.io/en/latest/).
+
+## Citation
+
+The framework is described in **PQuantML: A Tool for End-to-End Hardware-aware Model Compression** ([arXiv:2603.26595](https://arxiv.org/abs/2603.26595)).
+
+If you use PQuantML in your work, please cite:
+
+```bibtex
+@article{niemi2026pquantml,
+  title   = {PQuantML: A Tool for End-to-End Hardware-aware Model Compression},
+  author  = {Niemi, Roope and Petrovych, Anastasiia and Das, Arghya and
+             Lupi, Enrico and Sun, Chang and Danopoulos, Dimitrios and
+             Helbing, Marlon Joshua and Liu, Mia and Kagan, Michael and
+             Loncar, Vladimir and Pierini, Maurizio},
+  journal = {arXiv preprint arXiv:2603.26595},
+  year    = {2026}
+}
+```
+
+## Authors
+
+- Roope Niemi (CERN)
+- Anastasiia Petrovych (CERN)
+- Arghya Das (Purdue University)
+- Enrico Lupi (CERN)
+- Chang Sun (Caltech)
+- Dimitrios Danopoulos (CERN)
+- Marlon Joshua Helbing
+- Mia Liu (Purdue University)
+- Michael Kagan (SLAC National Accelerator Laboratory)
+- Vladimir Loncar (CERN)
+- Maurizio Pierini (CERN)

docs/Makefile

@@ -18,4 +18,4 @@ help:
 # "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
 %: Makefile
 	@sphinx-apidoc -f -o autodoc/ ../src/HGQ
-	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

docs/requirements.txt

@@ -1,5 +1,5 @@
-sphinx
 furo
 myst-parser
-sphinx_rtd_theme
+sphinx
 sphinx-autodoc-typehints
+sphinx_rtd_theme