sparsepixels 0.2.3__tar.gz → 0.3.0__tar.gz

{sparsepixels-0.2.3/sparsepixels.egg-info → sparsepixels-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sparsepixels
-Version: 0.2.3
+Version: 0.3.0
 Summary: Efficient convolution for sparse data on FPGAs
 Home-page: https://github.com/hftsoi/sparse-pixels
 Author: Ho Fung Tsoi
@@ -14,6 +14,7 @@ License-File: LICENSE
 Requires-Dist: tensorflow
 Requires-Dist: keras>=3.0
 Requires-Dist: HGQ2>=0.1.8
+Requires-Dist: matplotlib
 Dynamic: license-file
 
 <p align="center">
@@ -34,7 +35,7 @@ Dynamic: license-file
 [![arXiv](https://img.shields.io/badge/arXiv-2512.06208-b31b1b.svg?style=flat-square)](https://arxiv.org/abs/2512.06208)
 [![PyPI - Version](https://img.shields.io/pypi/v/sparsepixels?color=orange&style=flat-square)](https://pypi.org/project/sparsepixels)
 
-> **Note:** We are actively working on hls4ml integration to auto-convert sparse models to HLS, along with a major upgrade with partial parallelization and streaming for sparse layers in HLS. Stay tuned!
+SparsePixels is a Keras 3 library to build, train, and deploy sparse convolutional neural networks on FPGAs. In many detectors, especially in high-energy physics experiments, the images are almost empty: only a handful of pixels carry a signal (the hits), yet a standard CNN still spends compute on every pixel. A sparse CNN convolves only over the active pixels, so its cost scales with the number of hits rather than the image size, which is what makes low-latency, real-time inference (for example in a trigger) feasible on an FPGA. This library builds quantization-aware (via [HGQ2](https://github.com/calad0i/HGQ2)) sparse CNNs in which the pixel budget and the activity threshold can be learned from data, with a hardware-aware penalty that drives the budget toward the fewest pixels the task tolerates. Trained models convert to FPGA firmware through the [hls4ml](https://github.com/fastmachinelearning/hls4ml) integration, with control over the parallelization of the sparse layers to trade latency against resource usage.
 
 ## Installation
 
@@ -46,7 +47,7 @@ pip install sparsepixels
 
 ## Getting Started
 
-Import sparse layers and quantization library (HGQ2):
+Import the sparse layers, the quantization library (HGQ2), and the training utilities:
 
 ```python
 import keras
@@ -54,12 +55,28 @@ from keras.layers import Flatten, Activation
 from hgq.layers import QConv2D, QDense
 from hgq.config import QuantizerConfigScope, LayerConfigScope
 from hgq.quantizer.config import QuantizerConfig
-from sparsepixels.layers import InputReduce, QConv2DSparse, AveragePooling2DSparse
+from sparsepixels.layers import InputReduce, QConv2DSparse, AveragePooling2DSparse, MaxPooling2DSparse
+from sparsepixels.utils import (
+    active_pixels_vs_threshold, plot_reduced_examples,
+    set_sparse_ebops_factor, cosine_lr,
+    SparseTrainingMonitor, plot_history,
+    print_quantization, plot_quantization,
+)
+```
+
+First, study the data to pick a threshold and an initial pixel budget `n`: how many pixels stay
+active as the threshold rises, and what a candidate `(n, threshold)` keeps on a few images.
+
+```python
+active_pixels_vs_threshold(x_train)
+plot_reduced_examples(x_train, n=20, threshold=0.1, n_examples=4)
 ```
 
-Build an example sparse CNN within HGQ2 quantization scopes. A custom input quantizer
-config with higher initial fractional bits (`f0=8`) is used to prevent the default (`f0=2`)
-from zeroing out sparse signals in early training epochs:
+Build an example sparse CNN within HGQ2 quantization scopes. A custom input quantizer config with
+higher initial fractional bits (`f0=8`) prevents the default (`f0=2`) from zeroing out sparse signals
+in early training epochs. `InputReduce` keeps the first `n` active pixels (first channel above
+`threshold`); by default `n` and `threshold` are trainable hyperparameters, and a penalty of weight `beta_n`
+nudges the budget smaller, trading a little accuracy for lower FPGA latency and resources.
 
 ```python
 iq_conf = QuantizerConfig(place='datalane', q_type='kif', i0=4, f0=8, overflow_mode='WRAP')
@@ -71,8 +88,15 @@ with (
 ):
     x_in = keras.Input(shape=(28, 28, 1), name='x_in')
 
-    # Sparse input reduction: retain up to n_max_pixels active pixels
-    x, keep_mask = InputReduce(n_max_pixels=20, threshold=0.1, name='input_reduce')(x_in)
+    # Sparse input reduction
+    x, keep_mask = InputReduce(
+        n=30,                    # initial pixel budget
+        threshold=0.1,           # initial activity threshold
+        beta_n=1e-5,             # weight of the pixel budget penalty
+        learn_n=True,            # trainable pixel budget
+        learn_threshold=True,    # trainable threshold
+        name='input_reduce',
+    )(x_in)
 
     # Sparse convolution
     x = QConv2DSparse(filters=3, kernel_size=3, name='conv1', padding='same', strides=1,
@@ -88,6 +112,33 @@ with (
 model = keras.Model(x_in, x)
 ```
 
+Train the model, then read out the learned sparsity to deploy. `set_sparse_ebops_factor` makes the
+EBOPS (a proxy for the quantized hardware cost) reflect the sparse compute rather than a dense one; a
+cosine-decayed learning rate together with `restore_best_weights` keeps the learned budget from
+over-compressing near the end of training. `plot_history` shows the loss breakdown, the learned
+budget/threshold and the EBOPS in one figure, and the values to deploy are `layer.n_max_pixels` and
+`layer.threshold`.
+
+```python
+set_sparse_ebops_factor(model)
+
+steps_per_epoch = len(x_train) // 128
+early_stop = keras.callbacks.EarlyStopping(monitor='val_accuracy', mode='max', patience=20, restore_best_weights=True)
+model.compile(
+    optimizer=keras.optimizers.Adam(cosine_lr(1e-3, epochs=100, steps_per_epoch=steps_per_epoch)),
+    loss='categorical_crossentropy', metrics=['accuracy'],
+)
+history = model.fit(x_train, y_train, validation_data=(x_val, y_val),
+                    epochs=100, batch_size=128, callbacks=[early_stop, SparseTrainingMonitor()])
+
+plot_history(history, early_stopping=early_stop)   # loss breakdown, budget, threshold, EBOPS
+print_quantization(model)                          # per-layer bit-width distribution and EBOPS
+plot_quantization(model)
+
+ir = model.get_layer('input_reduce')
+print(f"deploy with n_max_pixels={ir.n_max_pixels}, threshold={ir.threshold:.3f}")
+```
+
 ## Converting a trained model to HLS with hls4ml
 
 > **Note:** A [PR](https://github.com/fastmachinelearning/hls4ml/pull/1468) adding `sparsepixels` support to the official [hls4ml](https://github.com/fastmachinelearning/hls4ml) repo has been submitted but is not yet merged. In the meantime you can install hls4ml from the PR branch on this fork to try the converter:
@@ -109,17 +160,17 @@ hls_model = hls4ml.converters.convert_from_keras_model(
     hls_config=hls_config,
     output_dir='hls_proj/my_sparse_cnn',
     backend='Vitis',
-    io_type='io_parallel',  # io_stream is not supported yet
+    io_type='io_parallel',
 )
 hls_model.write()
 hls_model.compile()
 y_hls = hls_model.predict(x_test)
 ```
 
-> **Note:** The converter currently supports only fully parallelized `io_parallel` HLS. We are working on expanding to partial parallelization and `io_stream` for larger flexibility.
-
 ## Documentation
 
+Coming soon!
+
 ## Citation
 
 If you find this useful in your research, please consider citing:
@@ -135,4 +186,3 @@ If you find this useful in your research, please consider citing:
     year = "2025"
 }
 ```
-

sparsepixels-0.3.0/README.md

@@ -0,0 +1,169 @@
+<p align="center">
+  <img src="https://raw.githubusercontent.com/hftsoi/sparse-pixels/main/docs/figs/logo.png" width="300" />
+</p>
+
+<p align="center">
+  <img src="https://raw.githubusercontent.com/hftsoi/sparse-pixels/main/docs/figs/sparsepixels.png" width="900"/>
+</p>
+
+<p align="center">
+  <img src="https://raw.githubusercontent.com/hftsoi/sparse-pixels/main/docs/figs/cnn_standard.gif" width="400" />
+  <img src="https://raw.githubusercontent.com/hftsoi/sparse-pixels/main/docs/figs/cnn_sparse.gif" width="400" />
+</p>
+
+# SparsePixels: Efficient convolution for sparse data on FPGAs
+
+[![arXiv](https://img.shields.io/badge/arXiv-2512.06208-b31b1b.svg?style=flat-square)](https://arxiv.org/abs/2512.06208)
+[![PyPI - Version](https://img.shields.io/pypi/v/sparsepixels?color=orange&style=flat-square)](https://pypi.org/project/sparsepixels)
+
+SparsePixels is a Keras 3 library to build, train, and deploy sparse convolutional neural networks on FPGAs. In many detectors, especially in high-energy physics experiments, the images are almost empty: only a handful of pixels carry a signal (the hits), yet a standard CNN still spends compute on every pixel. A sparse CNN convolves only over the active pixels, so its cost scales with the number of hits rather than the image size, which is what makes low-latency, real-time inference (for example in a trigger) feasible on an FPGA. This library builds quantization-aware (via [HGQ2](https://github.com/calad0i/HGQ2)) sparse CNNs in which the pixel budget and the activity threshold can be learned from data, with a hardware-aware penalty that drives the budget toward the fewest pixels the task tolerates. Trained models convert to FPGA firmware through the [hls4ml](https://github.com/fastmachinelearning/hls4ml) integration, with control over the parallelization of the sparse layers to trade latency against resource usage.
+
+## Installation
+
+With Python >= 3.10:
+
+```
+pip install sparsepixels
+```
+
+## Getting Started
+
+Import the sparse layers, the quantization library (HGQ2), and the training utilities:
+
+```python
+import keras
+from keras.layers import Flatten, Activation
+from hgq.layers import QConv2D, QDense
+from hgq.config import QuantizerConfigScope, LayerConfigScope
+from hgq.quantizer.config import QuantizerConfig
+from sparsepixels.layers import InputReduce, QConv2DSparse, AveragePooling2DSparse, MaxPooling2DSparse
+from sparsepixels.utils import (
+    active_pixels_vs_threshold, plot_reduced_examples,
+    set_sparse_ebops_factor, cosine_lr,
+    SparseTrainingMonitor, plot_history,
+    print_quantization, plot_quantization,
+)
+```
+
+First, study the data to pick a threshold and an initial pixel budget `n`: how many pixels stay
+active as the threshold rises, and what a candidate `(n, threshold)` keeps on a few images.
+
+```python
+active_pixels_vs_threshold(x_train)
+plot_reduced_examples(x_train, n=20, threshold=0.1, n_examples=4)
+```
+
+Build an example sparse CNN within HGQ2 quantization scopes. A custom input quantizer config with
+higher initial fractional bits (`f0=8`) prevents the default (`f0=2`) from zeroing out sparse signals
+in early training epochs. `InputReduce` keeps the first `n` active pixels (first channel above
+`threshold`); by default `n` and `threshold` are trainable hyperparameters, and a penalty of weight `beta_n`
+nudges the budget smaller, trading a little accuracy for lower FPGA latency and resources.
+
+```python
+iq_conf = QuantizerConfig(place='datalane', q_type='kif', i0=4, f0=8, overflow_mode='WRAP')
+
+with (
+    QuantizerConfigScope(place='all', default_q_type='kbi', overflow_mode='SAT_SYM'),
+    QuantizerConfigScope(place='datalane', default_q_type='kif', overflow_mode='WRAP'),
+    LayerConfigScope(enable_ebops=True, enable_iq=True, beta0=1e-5),
+):
+    x_in = keras.Input(shape=(28, 28, 1), name='x_in')
+
+    # Sparse input reduction
+    x, keep_mask = InputReduce(
+        n=30,                    # initial pixel budget
+        threshold=0.1,           # initial activity threshold
+        beta_n=1e-5,             # weight of the pixel budget penalty
+        learn_n=True,            # trainable pixel budget
+        learn_threshold=True,    # trainable threshold
+        name='input_reduce',
+    )(x_in)
+
+    # Sparse convolution
+    x = QConv2DSparse(filters=3, kernel_size=3, name='conv1', padding='same', strides=1,
+                      activation='relu', iq_conf=iq_conf)([x, keep_mask])
+
+    # Sparse pooling
+    x, keep_mask = AveragePooling2DSparse(2, name='pool1')([x, keep_mask])
+
+    x = Flatten(name='flatten')(x)
+    x = QDense(10, name='dense1', activation='relu', iq_conf=iq_conf)(x)
+    x = Activation('softmax', name='softmax')(x)
+
+model = keras.Model(x_in, x)
+```
+
+Train the model, then read out the learned sparsity to deploy. `set_sparse_ebops_factor` makes the
+EBOPS (a proxy for the quantized hardware cost) reflect the sparse compute rather than a dense one; a
+cosine-decayed learning rate together with `restore_best_weights` keeps the learned budget from
+over-compressing near the end of training. `plot_history` shows the loss breakdown, the learned
+budget/threshold and the EBOPS in one figure, and the values to deploy are `layer.n_max_pixels` and
+`layer.threshold`.
+
+```python
+set_sparse_ebops_factor(model)
+
+steps_per_epoch = len(x_train) // 128
+early_stop = keras.callbacks.EarlyStopping(monitor='val_accuracy', mode='max', patience=20, restore_best_weights=True)
+model.compile(
+    optimizer=keras.optimizers.Adam(cosine_lr(1e-3, epochs=100, steps_per_epoch=steps_per_epoch)),
+    loss='categorical_crossentropy', metrics=['accuracy'],
+)
+history = model.fit(x_train, y_train, validation_data=(x_val, y_val),
+                    epochs=100, batch_size=128, callbacks=[early_stop, SparseTrainingMonitor()])
+
+plot_history(history, early_stopping=early_stop)   # loss breakdown, budget, threshold, EBOPS
+print_quantization(model)                          # per-layer bit-width distribution and EBOPS
+plot_quantization(model)
+
+ir = model.get_layer('input_reduce')
+print(f"deploy with n_max_pixels={ir.n_max_pixels}, threshold={ir.threshold:.3f}")
+```
+
+## Converting a trained model to HLS with hls4ml
+
+> **Note:** A [PR](https://github.com/fastmachinelearning/hls4ml/pull/1468) adding `sparsepixels` support to the official [hls4ml](https://github.com/fastmachinelearning/hls4ml) repo has been submitted but is not yet merged. In the meantime you can install hls4ml from the PR branch on this fork to try the converter:
+>
+> ```bash
+> pip install "git+https://github.com/hftsoi/hls4ml.git@sparsepixels"
+> ```
+
+Once installed, converting a trained sparsepixels model to HLS is as usual:
+
+```python
+import hls4ml
+
+hls_config = hls4ml.utils.config_from_keras_model(model, granularity='name')
+hls_config.setdefault('Model', {})['PipelineStyle'] = 'dataflow'  # use "#pragma HLS DATAFLOW" (instead of the default "#pragma HLS PIPELINE" for io_parallel)
+
+hls_model = hls4ml.converters.convert_from_keras_model(
+    model,
+    hls_config=hls_config,
+    output_dir='hls_proj/my_sparse_cnn',
+    backend='Vitis',
+    io_type='io_parallel',
+)
+hls_model.write()
+hls_model.compile()
+y_hls = hls_model.predict(x_test)
+```
+
+## Documentation
+
+Coming soon!
+
+## Citation
+
+If you find this useful in your research, please consider citing:
+
+```
+@article{Tsoi:2025nvg,
+    author = "Tsoi, Ho Fung and Rankin, Dylan and Loncar, Vladimir and Harris, Philip",
+    title = "{SparsePixels: Efficient Convolution for Sparse Data on FPGAs}",
+    eprint = "2512.06208",
+    archivePrefix = "arXiv",
+    primaryClass = "cs.AR",
+    month = "12",
+    year = "2025"
+}
+```

{sparsepixels-0.2.3 → sparsepixels-0.3.0}/setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = sparsepixels
-version = 0.2.3
+version = 0.3.0
 description = Efficient convolution for sparse data on FPGAs
 author = Ho Fung Tsoi
 author_email = ho.fung.tsoi@cern.ch
@@ -20,6 +20,7 @@ install_requires =
 	tensorflow
 	keras>=3.0
 	HGQ2>=0.1.8
+	matplotlib
 include_package_data = True
 
 [options.package_data]

sparsepixels-0.3.0/sparsepixels/layers.py

@@ -0,0 +1,325 @@
+import keras
+from hgq.layers import QConv2D
+from hgq.quantizer import Quantizer
+from hgq.quantizer.config import QuantizerConfig
+from keras import ops
+from keras.layers import AveragePooling2D, MaxPooling2D
+
+
+class _ClipToRange(keras.constraints.Constraint):
+    """Weight constraint that clips values to the range [lo, hi] after each update."""
+
+    def __init__(self, lo, hi):
+        self.lo = float(lo)
+        self.hi = float(hi)
+
+    def __call__(self, w):
+        return ops.clip(w, self.lo, self.hi)
+
+    def get_config(self):
+        return {"lo": self.lo, "hi": self.hi}
+
+
+class InputReduce(keras.layers.Layer):
+    """Reduce a dense image to its first n active pixels for sparse FPGA inference.
+
+    Keeps the first n pixels whose first channel is above threshold, in raster order, and zeroes the
+    rest, returning the masked image together with a 0/1 keep mask that the following sparse layers
+    use as the sparse representation.
+
+    The budget n and the threshold can be learned during training (the default) so they need not be
+    tuned by hand; set learn_n or learn_threshold to False to keep either fixed (both False gives the
+    plain, non-learnable selection). The selection is always exact -- the learnable versions only
+    shape the gradient, so the layer behaves identically at inference and stays deployable. When n is
+    learned, a penalty of weight beta_n nudges it smaller, trading a little accuracy for lower FPGA
+    latency and resources; it starts at n and stays within [1, 4*n]. After training, read the values
+    to deploy from the n_max_pixels and threshold properties.
+
+    Args:
+        n: initial pixel budget, and the fixed budget when learn_n is False.
+        threshold: initial activity threshold on the first channel, fixed when learn_threshold is False.
+        beta_n: weight of the budget penalty added to the loss (0 disables it).
+        learn_n: make the pixel budget trainable.
+        learn_threshold: make the threshold trainable.
+        tau_threshold: softness of the threshold surrogate used to obtain gradients.
+        tau_n: softness of the budget-cutoff surrogate used to obtain gradients.
+    """
+
+    def __init__(
+        self,
+        n=30,
+        threshold=0.0,
+        beta_n=1e-5,
+        learn_n=True,
+        learn_threshold=True,
+        tau_threshold=0.05,
+        tau_n=1.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.n_init = int(n)
+        self.threshold_init = float(threshold)
+        self.beta_n = float(beta_n)
+        self.learn_n = learn_n
+        self.learn_threshold = learn_threshold
+        self.tau_threshold = float(tau_threshold)
+        self.tau_n = float(tau_n)
+
+    def build(self, input_shape):
+        if self.learn_threshold:
+            self.threshold_w = self.add_weight(
+                name="threshold",
+                shape=(),
+                initializer=keras.initializers.Constant(self.threshold_init),
+                trainable=True,
+                constraint=keras.constraints.NonNeg(),
+            )
+        if self.learn_n:
+            # Parametrize the budget as a fraction of its initial value (n = n_init * n_frac) so it
+            # moves at a useful rate under Adam; clip to [1, 4x] so it can shrink or modestly grow.
+            self.n_frac = self.add_weight(
+                name="n_frac",
+                shape=(),
+                initializer=keras.initializers.Constant(1.0),
+                trainable=True,
+                constraint=_ClipToRange(1.0 / self.n_init, 4.0),
+            )
+        super().build(input_shape)
+
+    def call(self, inputs):
+        dt = inputs.dtype
+        batch_size = ops.shape(inputs)[0]
+        h = ops.shape(inputs)[1]
+        w = ops.shape(inputs)[2]
+        score = ops.reshape(inputs[..., 0], [batch_size, h * w])
+
+        thr = self.threshold_w if self.learn_threshold else ops.cast(self.threshold_init, dt)
+        n = ops.cast(self.n_init, dt) * self.n_frac if self.learn_n else ops.cast(self.n_init, dt)
+
+        # Exact selection used for the forward pass.
+        active_hard = ops.cast(score > thr, dt)
+        rank_hard = ops.cumsum(active_hard, axis=1)
+        keep_hard = active_hard * ops.cast(rank_hard <= ops.round(n), dt)
+
+        if self.learn_threshold or self.learn_n:
+            # Differentiable surrogate used only for the gradient (straight-through to the exact
+            # selection above); the score > 0 gate keeps zero background pixels out when threshold=0.
+            active_soft = ops.sigmoid((score - thr) / self.tau_threshold) * ops.cast(score > 0, dt)
+            rank_soft = ops.cumsum(active_soft, axis=1)
+            keep_soft = active_soft * ops.sigmoid((n - rank_soft) / self.tau_n)
+            keep_flat = keep_soft + ops.stop_gradient(keep_hard - keep_soft)
+        else:
+            keep_flat = keep_hard
+
+        keep_mask = ops.reshape(keep_flat, [batch_size, h, w, 1])
+        inputs_reduced = inputs * keep_mask
+
+        if self.learn_n:
+            self.add_loss(self.beta_n * n)
+
+        return inputs_reduced, keep_mask
+
+    @property
+    def n_max_pixels(self):
+        """Integer pixel budget to deploy (the initial value until the layer is built)."""
+        if self.learn_n and self.built:
+            return int(round(self.n_init * float(ops.convert_to_numpy(self.n_frac))))
+        return self.n_init
+
+    @property
+    def threshold(self):
+        """Threshold to deploy (the initial value until the layer is built)."""
+        if self.learn_threshold and self.built:
+            return float(ops.convert_to_numpy(self.threshold_w))
+        return self.threshold_init
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "n": self.n_init,
+                "threshold": self.threshold_init,
+                "beta_n": self.beta_n,
+                "learn_n": self.learn_n,
+                "learn_threshold": self.learn_threshold,
+                "tau_threshold": self.tau_threshold,
+                "tau_n": self.tau_n,
+            }
+        )
+        return config
+
+
+class RemoveDilatedPixels(keras.layers.Layer):
+    """Re-apply the keep mask, zeroing every pixel that is not active.
+
+    Multiplies a feature map by its 0/1 keep mask (broadcast over channels) so only the kept pixels
+    carry values. Used inside the sparse layers to restore the sparse representation after a dense op.
+
+    Call args:
+        inputs: tuple (x, mask) of the feature map and its keep mask.
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def call(self, inputs):
+        x, mask = inputs
+        mask = ops.cast(mask, x.dtype)
+        return x * mask
+
+    def get_config(self):
+        return super().get_config()
+
+
+class QConv2DSparse(keras.layers.Layer):
+    """Quantized 2D convolution that operates on the sparse (active-pixel) representation.
+
+    Wraps an HGQ QConv2D: masks the input to the active pixels, convolves, adds a separately
+    quantized per-filter bias on the nonzero outputs, applies the activation, then re-masks the
+    output. This is numerically the same as a dense quantized conv restricted to the active pixels,
+    which is what the HLS sparse_conv kernel computes.
+
+    Args:
+        *conv_args: positional arguments forwarded to hgq.layers.QConv2D (e.g. filters, kernel_size).
+        **conv_kwargs: keyword arguments forwarded to QConv2D (padding, strides, ...). use_bias,
+            activation and bq_conf are handled here: the bias has its own weight and quantizer
+            (bq_conf), and the activation is applied after the bias.
+
+    Call args:
+        inputs: tuple (x, keep_mask) of the feature map and its keep mask.
+    """
+
+    def __init__(self, *conv_args, **conv_kwargs):
+        super().__init__(name=conv_kwargs.get("name", None))
+        self._use_bias = conv_kwargs.pop("use_bias", True)
+        self._bq_conf = conv_kwargs.pop("bq_conf", None) or QuantizerConfig("default", "bias")
+        self._activation = keras.activations.get(conv_kwargs.pop("activation", None))
+
+        conv_kwargs["use_bias"] = False
+        conv_kwargs["activation"] = None
+        self.conv = QConv2D(*conv_args, **conv_kwargs)
+        self.masker = RemoveDilatedPixels()
+
+    def build(self, input_shape):
+        # Build the wrapped conv eagerly here rather than lazily in call(): building it while Keras
+        # symbolically traces call() triggers an HGQ weight check that fails in graph mode.
+        x_shape = input_shape[0]
+        if not self.conv.built:
+            self.conv.build(x_shape)
+        if self._use_bias:
+            self.sparse_bias = self.add_weight(
+                name="sparse_bias",
+                shape=(self.conv.filters,),
+                initializer="zeros",
+                trainable=True,
+            )
+            self._bq = Quantizer(self._bq_conf, name=f"{self.name}_bq")
+            self._bq.build((self.conv.filters,))
+        super().build(input_shape)
+
+    def compute_output_shape(self, input_shape):
+        # Return the shape directly so Keras does not trace call() (masking preserves the shape).
+        return self.conv.compute_output_shape(input_shape[0])
+
+    def call(self, inputs, **kwargs):
+        x, keep_mask = inputs
+        x = self.masker((x, keep_mask))
+        y = self.conv(x, **kwargs)
+
+        if self._use_bias:
+            b = self._bq(self.sparse_bias)
+            b = ops.reshape(b, (1, 1, 1, -1))
+            non_zero = ops.cast(y != 0, y.dtype)
+            y = y + b * non_zero
+
+        if self._activation is not None:
+            y = self._activation(y)
+
+        y = self.masker((y, keep_mask))
+        return y
+
+    def get_config(self):
+        cfg = super().get_config()
+        cfg["conv_config"] = self.conv.get_config()
+        cfg["use_bias"] = self._use_bias
+        cfg["bq_conf"] = self._bq_conf
+        cfg["activation"] = keras.activations.serialize(self._activation)
+        return cfg
+
+    @classmethod
+    def from_config(cls, config):
+        conv_cfg = config.pop("conv_config")
+        use_bias = config.pop("use_bias", True)
+        bq_conf = config.pop("bq_conf", None)
+        activation = config.pop("activation", None)
+        return cls(**conv_cfg, use_bias=use_bias, bq_conf=bq_conf, activation=activation)
+
+
+class AveragePooling2DSparse(keras.layers.Layer):
+    """Average pooling on the sparse representation.
+
+    Average-pools the feature map and max-pools the keep mask, so a pooled cell stays active when any
+    of its source pixels were active. Mirrors the HLS sparse_pooling_avg kernel.
+
+    Args:
+        *pool_args: positional arguments forwarded to keras AveragePooling2D (e.g. pool_size).
+        **pool_kwargs: keyword arguments forwarded to the pooling layers.
+
+    Call args:
+        inputs: tuple (x, keep_mask) of the feature map and its keep mask.
+    """
+
+    def __init__(self, *pool_args, **pool_kwargs):
+        super().__init__(name=pool_kwargs.get("name", None))
+        self.avg_pool = AveragePooling2D(*pool_args, **pool_kwargs)
+        self.max_pool = MaxPooling2D(*pool_args, **pool_kwargs)
+
+    def call(self, inputs, **kwargs):
+        x, keep_mask = inputs
+        y = self.avg_pool(x, **kwargs)
+        keep_mask_pooled = self.max_pool(keep_mask)
+        return y, keep_mask_pooled
+
+    def get_config(self):
+        cfg = super().get_config()
+        cfg["pool_config"] = self.avg_pool.get_config()
+        return cfg
+
+    @classmethod
+    def from_config(cls, config):
+        pool_cfg = config.pop("pool_config")
+        return cls(**pool_cfg)
+
+
+class MaxPooling2DSparse(keras.layers.Layer):
+    """Max pooling on the sparse representation.
+
+    Max-pools both the feature map and the keep mask. Mirrors the HLS sparse_pooling_max kernel.
+
+    Args:
+        *pool_args: positional arguments forwarded to keras MaxPooling2D (e.g. pool_size).
+        **pool_kwargs: keyword arguments forwarded to the pooling layer.
+
+    Call args:
+        inputs: tuple (x, keep_mask) of the feature map and its keep mask.
+    """
+
+    def __init__(self, *pool_args, **pool_kwargs):
+        super().__init__(name=pool_kwargs.get("name", None))
+        self.max_pool = MaxPooling2D(*pool_args, **pool_kwargs)
+
+    def call(self, inputs, **kwargs):
+        x, keep_mask = inputs
+        y = self.max_pool(x, **kwargs)
+        keep_mask_pooled = self.max_pool(keep_mask)
+        return y, keep_mask_pooled
+
+    def get_config(self):
+        cfg = super().get_config()
+        cfg["pool_config"] = self.max_pool.get_config()
+        return cfg
+
+    @classmethod
+    def from_config(cls, config):
+        pool_cfg = config.pop("pool_config")
+        return cls(**pool_cfg)