LeNet-5 Subsample Layer in Tensorflow In Tensorflow, how do you implement the LeNet-5 pooling layers with trainable coefficient and bias terms?
Reading through the LeNet-5 paper, the subsample layers are described as follows:

Layer S2 is a sub-sampling layer with 6 feature maps of size 14x14. Each unit in each feature map is connected to a 2x2 neighborhood in the corresponding feature map in C1. The fout inputs to a unit in S2 are added, then multiplied by a trainable coefficient, and added to a trainable bias. The result is passed through a sigmoidal function. The 2x2 receptive fields are non-overlapping, therefore feature maps in S2 have half the number of rows and columns as feature maps in C1. Layer S2 has 12 trainable parameters and 5,880 connections.

http://yann.lecun.com/exdb/publis/pdf/lecun-98.pdf
However, in my search for examples of implementing LeNet-5 in Tensorflow, I haven't seen this pooling layer implemented with the trainable coefficient and bias. Instead, something like the following is used:
model = keras.Sequential()
model.add(layers.Conv2D(filters=6, 
                        kernel_size=(5, 5), 
                        activation='tanh', 
                        input_shape=(28,28,1),
                        padding='same'))
model.add(layers.AveragePooling2D(pool_size=(2, 2), 
                                  strides=(2, 2), 
                                  padding='valid'))
model.add(layers.Conv2D(filters=16, 
                        kernel_size=(5, 5), 
                        activation='tanh',
                        padding='valid'))
model.add(layers.AveragePooling2D(pool_size=(2, 2), 
                                  strides=(2, 2), 
                                  padding='valid'))
model.add(layers.Flatten())
model.add(layers.Dense(units=120, activation='tanh'))
model.add(layers.Dense(units=84, activation='tanh'))
model.add(layers.Dense(units=10, activation = 'softmax'))

Calling model.summary() on a model like this yields:
Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
conv2d (Conv2D)              (None, 28, 28, 6)         156       
_________________________________________________________________
average_pooling2d (AveragePo (None, 14, 14, 6)         0         
_________________________________________________________________
conv2d_1 (Conv2D)            (None, 10, 10, 16)        2416      
_________________________________________________________________
average_pooling2d_1 (Average (None, 5, 5, 16)          0         
_________________________________________________________________
flatten (Flatten)            (None, 400)               0         
_________________________________________________________________
dense (Dense)                (None, 120)               48120     
_________________________________________________________________
dense_1 (Dense)              (None, 84)                10164     
_________________________________________________________________
dense_2 (Dense)              (None, 10)                850       
=================================================================
Total params: 61,706
Trainable params: 61,706
Non-trainable params: 0
_________________________________________________________________

The pooling layers have no trainable parameters. Maybe those parameters aren't so important for performance, but I'm curious how to implement the original pooling layers in Tensorflow.
 A: I similarly found that most folks use either AveragePooling2D or MaxPooling2D because there doesn't appear to exist a layer in TensorFlow or Keras that precisely matches the Subsampling layer as described in the LeNet-5 paper, so I ended up implementing it myself. The number of trainable parameters matches the paper's description. Please let me know if you have any suggestions for how to improve it, or if you see any bugs.
Full code for my Subsampling layer is below; if you want to see how I'm using it, please see my LeNet notebooks — be sure to select one of the versions that use this Subsampling layer; I am maintaining multiple implementations for comparison purposes, and some of them use MaxPooling2D instead.
In my experience, MaxPooling2D actually seems to do a little bit better than the custom Subsampling layer, but there are a lot more details to that paper that most simple implementations overlook (custom activation function, pixel value scaling, training set augmentation, etc.), so it may actually be better to implement all of the details in the paper, but for a very simple implementation, MaxPooling2D seems to work just as well.
# Copyright 2022 Google LLC.
# SPDX-License-Identifier: Apache-2.0
"""Subsampling layer as defined in the LeNet paper."""

from typing import Callable, List, Optional, Tuple, Union

import numpy as np

import tensorflow as tf
from tensorflow import keras
from keras.layers import Layer


class SubsamplingArgumentError(ValueError):
    pass


class Subsampling(Layer):
    pool_size: Tuple[int, int]
    strides: Tuple[int, int]
    padding: str
    activation: Callable[[float], float]
    w: np.array
    b: np.array

    def __init__(
        self,
        pool_size: Union[int, List[int], Tuple[int, int]] = (2, 2),
        strides: Optional[Union[int, List[int], Tuple[int, int]]] = None,
        padding: str = 'VALID',
        activation: Callable[[float], float] = None,
        **kwargs):
        """Subsampling layer as described in the LeNet paper.
        This layer is not a simple average or max pooling that are typically
        used to implement LeNet: neither average-pooling nor max-pooling have
        any trainable parameters, while the subsampling layer described in the
        LeNet paper *does* have trainable parameters.
        Note: assumes data format is `(batch_size, rows, cols, channels)`, i.e.,
        what TensorFlow / Keras describe as "channels_last".
        Args:
          pool_size: int or 2-tuple specifying pool size (aka kernel size)
          strides: int or 2-tuple; if unspecified, will be copied from
            `pool_size`
          padding:
        """
        super().__init__(**kwargs)

        if isinstance(pool_size, int):
            self.pool_size = (pool_size, pool_size)
        elif (isinstance(pool_size, list) or isinstance(pool_size, tuple)) and \
            len(pool_size) == 2:
            self.pool_size = tuple(pool_size)
        else:
            raise SubsamplingArgumentError(
                f"`pool_size` must be an int or 2-tuple; received: {pool_size}")

        if strides is None:
            self.strides == self.pool_size
        elif isinstance(strides, int):
            self.strides = (strides, strides)
        elif (isinstance(strides, list) or isinstance(strides, tuple)) and \
            len(strides) == 2:
            self.strides = tuple(strides)
        else:
            raise SubsamplingArgumentError(
                f"`strides` must be an int or 2-tuple; received: {strides}")

        assert padding is not None and padding.upper() in ('VALID', 'SAME'), (
            f"`padding` must be either 'VALID' or 'SAME'; received: {padding}")
        self.padding = padding.upper()

        self.activation = activation or (lambda x: x)

    def build(
        self, input_shape: Union[List[Optional[int]],
                                 Tuple[Optional[int], int, int, int]]) -> None:
        """Builds internal structures to prepare for model training.
        Args:
          input_shape: length-4 list or tuple representing (batch_size, rows,
            cols, channels); where `batch_size` may be None; see docs above for
            `__init__()` for details.
        """
        if len(input_shape) != 4:
            raise SubsamplingArgumentError(
                f"`len(input_shape)` != 4; received: {input_shape}")
        if input_shape[0] is not None:
            raise SubsamplingArgumentError(
                f"`input_shape[0] must be None; received: {input_shape}")

        in_chan = input_shape[3]

        self.w = self.add_weight(shape=(1, 1, in_chan),
                                 initializer="random_normal",
                                 trainable=True)

        self.b = self.add_weight(shape=(1, 1, in_chan),
                                 initializer="random_normal",
                                 trainable=True)

    def call(self, inputs: tf.Tensor) -> tf.Tensor:
        """Computes subsampling value: `w * (sum of window entries) + b`."""

        # `scale` here undoes the average pooling by getting the original sum,
        # which is what we need, but there isn't a pooling mechanism that just
        # gets us the sum of products.
        scale = self.pool_size[0] * self.pool_size[1]
        tf_scale = tf.constant(scale, dtype='float32')

        avg = tf.nn.pool(inputs,
                         window_shape=self.pool_size,
                         pooling_type='AVG',
                         strides=self.strides,
                         padding=self.padding)

        return self.activation(self.w * tf_scale * avg + self.b)

Here's the output of model.summary() for comparison:
Model: "LeNet-5"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 C1 (Conv2D)                 (None, 28, 28, 6)         156       
                                                                 
 S2 (Subsampling)            (None, 14, 14, 6)         12        
                                                                 
 C3 (Conv2D)                 (None, 10, 10, 16)        2416      
                                                                 
 S4 (Subsampling)            (None, 5, 5, 16)          32        
                                                                 
 flatten (Flatten)           (None, 400)               0         
                                                                 
 C5 (Dense)                  (None, 120)               48120     
                                                                 
 F6 (Dense)                  (None, 84)                10164     
                                                                 
 Output (Dense)              (None, 10)                850       
                                                                 
=================================================================
Total params: 61,750
Trainable params: 61,750
Non-trainable params: 0
_________________________________________________________________

