Repository: marl/autopool (language: python)

Issue:
Restructuring for multiple backends
Tagging @justinsalamon for some input on this, should be easy.

The reference implementation of autopool is tied to old-style Keras, and that's cool.  PR #12 updates this to work with v2.2 and up.  However, it is recommended at this point that keras-dependent code target the tf.keras backend instead of the standalone package.

I think we should still provide the old-style interface (for historical reasons), but we should also have a tf.keras-facing implementation.

At the same time, we may as well provide a pytorch implementation as well.

I think the sensible way to approach this is by having different submodules for each backend, eg `autopool.keras`, `autopool.tf`, `autopool.torch` from which you would import the AutoPool/SoftMaxPool classes.

The only hiccup here is that the import structure changes.  Instead of the previous

```python
from autopool import AutoPool1D
```

you would have

```python
from autopool.keras import AutoPool1D
```
or
```python
from autopool.torch import AutoPool1D
```

(not unlike `from tqdm.notebook import tqdm`).

What do you think?


Discussion / hints:

Fix patch:
diff --git a/.github/environment-ci.yml b/.github/environment-ci.yml
new file mode 100644
index 0000000..d133bf9
--- /dev/null
+++ b/.github/environment-ci.yml
@@ -0,0 +1,12 @@
+name: test
+channels:
+  - conda-forge
+  - defaults
+dependencies:
+  # required
+  - keras>=2.2
+  - tensorflow>=2.0
+  - numpy >= 1.17
+  - pytest-cov
+  - pytest-faulthandler
+  - pytest
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
new file mode 100644
index 0000000..309d5c1
--- /dev/null
+++ b/.github/workflows/ci.yml
@@ -0,0 +1,86 @@
+name: CI
+
+on:
+  push:
+    branches:
+    - main
+  pull_request:
+    branches:
+    - main
+
+concurrency:
+    group: ${{ github.workflow }}-${{ github.ref }}
+    cancel-in-progress: True
+
+jobs:
+    test:
+        name: "Python ${{ matrix.python-version }} on ${{ matrix.os }}"
+        runs-on: ${{ matrix.os }}
+
+        strategy:
+            fail-fast: false
+            matrix:
+                include:
+                    - os: ubuntu-latest
+                      python-version: "3.7"
+                      channel-priority: "strict"
+                      envfile: ".github/environment-ci.yml"
+
+                    - os: ubuntu-latest
+                      python-version: "3.8"
+                      channel-priority: "strict"
+                      envfile: ".github/environment-ci.yml"
+
+                    - os: ubuntu-latest
+                      python-version: "3.9"
+                      channel-priority: "strict"
+                      envfile: ".github/environment-ci.yml"
+
+        steps:
+        - uses: actions/checkout@v2
+        - name: Cache conda
+          uses: actions/cache@v2
+          env:
+            # Increase this value to reset cache if etc/example-environment.yml has not changed
+            CACHE_NUMBER: 1
+          with:
+            path: ~/conda_pkgs_dir
+            key: ${{ runner.os }}-${{ matrix.python-version }}-conda-${{ env.CACHE_NUMBER }}-${{ hashFiles( matrix.envfile ) }}
+
+        - name: Install Conda environment
+          uses: conda-incubator/setup-miniconda@v2
+          with:
+            auto-update-conda: true
+            python-version: ${{ matrix.python-version }}
+            add-pip-as-python-dependency: true
+            auto-activate-base: false
+            activate-environment: test
+            channel-priority: ${{ matrix.channel-priority }}
+            environment-file: ${{ matrix.envfile }}
+            use-only-tar-bz2: true # IMPORTANT: This needs to be set for caching to work properly!
+
+        - name: Conda info
+          shell: bash -l {0}
+          run: |
+            conda info -a
+            conda list
+
+        - name: Install autopool
+          shell: bash -l {0}
+          run: python -m pip install --upgrade-strategy only-if-needed -e .[tests]
+
+        - name: Run pytest
+          shell: bash -l {0}
+          run: pytest -v
+
+        - name: Upload coverage to Codecov
+          uses: codecov/codecov-action@v3
+          with:
+            token: ${{ secrets.CODECOV_TOKEN }}
+            files: ./coverage.xml
+            directory: ./coverage/reports/
+            flags: unittests
+            env_vars: OS,PYTHON
+            name: codecov-umbrella
+            fail_ci_if_error: true
+            verbose: true
diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
new file mode 100644
index 0000000..3c7ddf3
--- /dev/null
+++ b/.github/workflows/publish.yml
@@ -0,0 +1,37 @@
+name: Publish Python 🐍 distributions 📦 to PyPI and TestPyPI
+
+on: 
+  release:
+    types: [created] 
+
+jobs:
+  build-n-publish:
+    name: Build and publish Python 🐍 distributions 📦 to PyPI and TestPyPI
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@master
+    - name: Set up Python 3.7
+      uses: actions/setup-python@v1
+      with:
+        python-version: 3.7
+
+    - name: Install pypa/build
+      run: >-
+        python -m
+        pip install
+        build
+        --user
+    - name: Build a binary wheel and a source tarball
+      run: >-
+        python -m
+        build
+        --sdist
+        --wheel
+        --outdir dist/
+        .
+    - name: Publish distribution 📦 to PyPI
+      if: startsWith(github.ref, 'refs/tags')
+      uses: pypa/gh-action-pypi-publish@master
+      with:
+        password: ${{ secrets.PYPI_API_TOKEN }}
diff --git a/LICENSE b/LICENSE.md
similarity index 100%
rename from LICENSE
rename to LICENSE.md
diff --git a/README.md b/README.md
index 8dcccd0..2f17518 100644
--- a/README.md
+++ b/README.md
@@ -3,7 +3,7 @@ Adaptive pooling operators for Multiple Instance Learning ([documentation](http:
 
 [![PyPI](https://img.shields.io/pypi/v/autopool.svg)](https://pypi.python.org/pypi/autopool)
 [![License](https://img.shields.io/github/license/mashape/apistatus.svg)](https://opensource.org/licenses/MIT)
-[![Documentation Status](https://readthedocs.org/projects/autopool/badge/?version=latest)](http://scaper.readthedocs.io/en/latest/?badge=latest)
+[![Documentation Status](https://readthedocs.org/projects/autopool/badge/?version=latest)](http://autopool.readthedocs.io/en/latest/?badge=latest)
 [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/autopool.svg)]()
 
 
@@ -19,8 +19,13 @@ IEEE Transactions on Audio, Speech and Language Processing, In press, 2018.
 Installation
 ------------
 
+To install the keras-based implementation:
 ```
-pip install autopool
+python -m pip install autopool[keras]
+```
+For the tensorflow implementation:
+```
+python -m pip install autopool[tf]
 ```
 
 Definition
@@ -37,7 +42,8 @@ Usage
 AutoPool is implemented as a [keras](https://keras.io/) layer, so using it is as straightforward as using any standard Keras pooling layer, for example:
 
 ```
-bag_pred = AutoPool(axis=1)(instance_pred)
+from autpool.keras import AutoPool1D
+bag_pred = AutoPool1D(axis=1)(instance_pred)
 ```
 
 Further details and examples are provided in the [documentation](http://autopool.readthedocs.io/).
@@ -48,23 +54,29 @@ Constrained and Regularized AutoPool
 In the [paper](http://www.justinsalamon.com/uploads/4/3/9/4/4394963/mcfee_autopool_taslp_2018.pdf) we show there can be benefits to either constraining the range α can take, or, alternatively, applying l2 regularization on α; this results in constrained AutoPool (CAP) and regularized AutoPool (RAP) respectively. Since AutoPool is implemented as a [keras](https://keras.io/) layer, CAP and RAP can be can be achieved through the layer's optional arugments:
 
 CAP with non-negative α:
-```
-bag_pred = AutoPool(axis=1, kernel_constraint=keras.constraints.non_neg())(instance_pred)
+```python
+bag_pred = AutoPool1D(axis=1, kernel_constraint=keras.constraints.non_neg())(instance_pred)
 ```
 
 CAP with α norm-constrained to some value `alpha_max`:
-```
-bag_pred = AutoPool(axis=1, kernel_constraint=keras.constraints.max_norm(alpha_max, axis=0))(instance_pred)
+```python
+bag_pred = AutoPool1D(axis=1, kernel_constraint=keras.constraints.max_norm(alpha_max, axis=0))(instance_pred)
 ```
 Heuristics for determining sensible values of `alpha_max` are given in the paper (section III.E).
 
 RAP with l2 regularized α:
-```
-bag_pred = AutoPool(axis=1, kernel_regularizer=keras.regularizers.l2(l=1e-4))(instance_pred)
+```python
+bag_pred = AutoPool1D(axis=1, kernel_regularizer=keras.regularizers.l2(l=1e-4))(instance_pred)
 ```
 
 CAP and RAP can be combined, of course, by applying both a kernel constraint and a kernel regularizer.
 
+If using the tensorflow-based implementation, all of the above will also work, except that `keras` should be replaced by `tensorflow.keras` (or `tf.keras`), e.g.:
+```python
+import tensorflow as tf
+from autopool.tf import AutoPool1D
+bag_pred = AutoPool1D(axis=1, kernel_regularizer=tf.keras.regularizers.l2(l=1e-4))(instance_pred)
+```
 
 Multi-label
 -----------
diff --git a/autopool/__init__.py b/autopool/__init__.py
index b4687ca..ff10bda 100644
--- a/autopool/__init__.py
+++ b/autopool/__init__.py
@@ -2,4 +2,3 @@
 '''AutoPool: adaptive pooling operators for multiple instance learning'''
 
 from .version import version as __version__
-from .autopool import *
diff --git a/autopool/autopool.py b/autopool/keras.py
similarity index 88%
rename from autopool/autopool.py
rename to autopool/keras.py
index 829a244..ca7a4a9 100644
--- a/autopool/autopool.py
+++ b/autopool/keras.py
@@ -2,14 +2,14 @@
 '''Autopool: Adaptive pooling operators for multiple instance learning'''
 
 from keras import backend as K
-from keras.engine.topology import Layer, InputSpec
+from keras.layers import Layer, InputSpec
 from keras import initializers
 from keras import constraints
 from keras import regularizers
 
 
 class AutoPool1D(Layer):
-    '''Automatically tuned soft-max pooling.
+    '''Automatically tuned soft-max pooling.  (Keras implementation)
 
     This layer automatically adapts the pooling behavior to interpolate
     between mean- and max-pooling for each dimension.
@@ -76,10 +76,7 @@ def get_config(self):
         return dict(list(base_config.items()) + list(config.items()))
 
     def call(self, x, mask=None):
-        scaled = self.kernel * x
-        max_val = K.max(scaled, axis=self.axis, keepdims=True)
-        softmax = K.exp(scaled - max_val)
-        weights = softmax / K.sum(softmax, axis=self.axis, keepdims=True)
+        weights = K.softmax(self.kernel * x, axis=self.axis)
         return K.sum(x * weights, axis=self.axis, keepdims=False)
 
 
@@ -111,9 +108,7 @@ def compute_output_shape(self, input_shape):
         return self.get_output_shape_for(input_shape)
 
     def call(self, x, mask=None):
-        max_val = K.max(x, axis=self.axis, keepdims=True)
-        softmax = K.exp((x - max_val))
-        weights = softmax / K.sum(softmax, axis=self.axis, keepdims=True)
+        weights = K.softmax(x, axis=self.axis)
         return K.sum(x * weights, axis=self.axis, keepdims=False)
 
     def get_config(self):
diff --git a/autopool/tf.py b/autopool/tf.py
new file mode 100644
index 0000000..0740217
--- /dev/null
+++ b/autopool/tf.py
@@ -0,0 +1,117 @@
+#!/usr/bin/env python
+'''Autopool: Adaptive pooling operators for multiple instance learning'''
+
+from tensorflow.keras import backend as K
+from tensorflow.keras.layers import Layer, InputSpec
+from tensorflow.keras import initializers
+from tensorflow.keras import constraints
+from tensorflow.keras import regularizers
+
+
+class AutoPool1D(Layer):
+    '''Automatically tuned soft-max pooling.  (tensorflow.keras implementation)
+
+    This layer automatically adapts the pooling behavior to interpolate
+    between mean- and max-pooling for each dimension.
+    '''
+    def __init__(self, axis=0,
+                 kernel_initializer='zeros',
+                 kernel_constraint=None,
+                 kernel_regularizer=None,
+                 **kwargs):
+        '''
+
+        Parameters
+        ----------
+        axis : int
+            Axis along which to perform the pooling. By default 0 (should be time).
+
+        kernel_initializer: Initializer for the weights matrix
+
+        kernel_regularizer: Regularizer function applied to the weights matrix
+
+        kernel_constraint: Constraint function applied to the weights matrix
+        kwargs
+        '''
+
+        if 'input_shape' not in kwargs and 'input_dim' in kwargs:
+            kwargs['input_shape'] = (kwargs.pop('input_dim'), )
+
+        super(AutoPool1D, self).__init__(**kwargs)
+
+        self.axis = axis
+        self.kernel_initializer = initializers.get(kernel_initializer)
+        self.kernel_constraint = constraints.get(kernel_constraint)
+        self.kernel_regularizer = regularizers.get(kernel_regularizer)
+        self.input_spec = InputSpec(min_ndim=3)
+        self.supports_masking = True
+
+    def build(self, input_shape):
+        assert len(input_shape) >= 3
+        input_dim = input_shape[-1]
+
+        self.kernel = self.add_weight(shape=(1, input_dim),
+                                      initializer=self.kernel_initializer,
+                                      name='kernel',
+                                      regularizer=self.kernel_regularizer,
+                                      constraint=self.kernel_constraint)
+        self.input_spec = InputSpec(min_ndim=2, axes={-1: input_dim})
+        self.built = True
+
+    def compute_output_shape(self, input_shape):
+        return self.get_output_shape_for(input_shape)
+
+    def get_output_shape_for(self, input_shape):
+        shape = list(input_shape)
+        del shape[self.axis]
+        return tuple(shape)
+
+    def get_config(self):
+        config = {'kernel_initializer': initializers.serialize(self.kernel_initializer),
+                  'kernel_constraint': constraints.serialize(self.kernel_constraint),
+                  'kernel_regularizer': regularizers.serialize(self.kernel_regularizer),
+                  'axis': self.axis}
+
+        base_config = super(AutoPool1D, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+    def call(self, x, mask=None):
+        weights = K.softmax(self.kernel * x, axis=self.axis)
+        return K.sum(x * weights, axis=self.axis, keepdims=False)
+
+
+class SoftMaxPool1D(Layer):
+    '''
+    Tensorflow-keras softmax pooling layer.
+    '''
+
+    def __init__(self, axis=0, **kwargs):
+        '''
+
+        Parameters
+        ----------
+        axis : int
+            Axis along which to perform the pooling. By default 0
+            (should be time).
+        kwargs
+        '''
+        super(SoftMaxPool1D, self).__init__(**kwargs)
+
+        self.axis = axis
+
+    def get_output_shape_for(self, input_shape):
+        shape = list(input_shape)
+        del shape[self.axis]
+        return tuple(shape)
+
+    def compute_output_shape(self, input_shape):
+        return self.get_output_shape_for(input_shape)
+
+    def call(self, x, mask=None):
+        weights = K.softmax(x, axis=self.axis)
+        return K.sum(x * weights, axis=self.axis, keepdims=False)
+
+    def get_config(self):
+        config = {'axis': self.axis}
+        base_config = super(SoftMaxPool1D, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
diff --git a/autopool/version.py b/autopool/version.py
index e1cd5a0..ac43fc5 100644
--- a/autopool/version.py
+++ b/autopool/version.py
@@ -2,5 +2,5 @@
 # -*- coding: utf-8 -*-
 """Version info"""
 
-short_version = '0.1.0'
-version = '0.1.0'
+short_version = '0.2.0'
+version = '0.2.0'
diff --git a/docs/api.rst b/docs/api.rst
index 68a9d71..c56c102 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -4,11 +4,25 @@
 
 API Reference
 -------------
-.. autoclass:: autopool.AutoPool1D
+
+Keras implementation
+^^^^^^^^^^^^^^^^^^^^
+.. autoclass:: autopool.keras.AutoPool1D
+    :no-undoc-members:
+    :special-members: __init__
+
+.. autoclass:: autopool.keras.SoftMaxPool1D
+    :no-undoc-members:
+    :special-members: __init__
+
+Tensorflow implementation
+^^^^^^^^^^^^^^^^^^^^^^^^^
+
+.. autoclass:: autopool.tf.AutoPool1D
     :no-undoc-members:
     :special-members: __init__
 
-.. autoclass:: autopool.SoftMaxPool1D
+.. autoclass:: autopool.tf.SoftMaxPool1D
     :no-undoc-members:
     :special-members: __init__
 
diff --git a/docs/conf.py b/docs/conf.py
index 5074c4b..191f460 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -75,7 +75,8 @@
 
 
 MOCK_MODULES = [
-    'keras', 'keras.engine', 'keras.engine.topology'
+    'keras', 'keras.layers',
+    'tensorflow.keras', 'tensorflow.keras.layers'
 ]
 
 MOCK_CLASSES = [
@@ -129,7 +130,7 @@ def __getattr__(cls, name):
 # Add any paths that contain custom static files (such as style sheets) here,
 # relative to this directory. They are copied after the builtin static files,
 # so a file named "default.css" will overwrite the builtin "default.css".
-html_static_path = ['_static']
+# html_static_path = ['_static']
 
 # Custom sidebar templates, must be a dictionary that maps document names
 # to template names.
diff --git a/docs/definitions.rst b/docs/definitions.rst
index 3a1d476..54ef89a 100644
--- a/docs/definitions.rst
+++ b/docs/definitions.rst
@@ -30,20 +30,20 @@ CAP with non-negative α:
 
 .. code-block:: python
 
-    bag_pred = AutoPool(axis=1, kernel_constraint=keras.constraints.non_neg())(instance_pred)
+    bag_pred = AutoPool1D(axis=1, kernel_constraint=keras.constraints.non_neg())(instance_pred)
 
 
 CAP with α norm-constrained to some value `alpha_max`:
 
 .. code-block:: python
 
-    bag_pred = AutoPool(axis=1, kernel_constraint=keras.constraints.max_norm(alpha_max, axis=0))(instance_pred)
+    bag_pred = AutoPool1D(axis=1, kernel_constraint=keras.constraints.max_norm(alpha_max, axis=0))(instance_pred)
 
 RAP with l2 regularized α:
 
 .. code-block:: python
 
-    bag_pred = AutoPool(axis=1, kernel_regularizer=keras.regularizers.l2(l=1e-4))(instance_pred)
+    bag_pred = AutoPool1D(axis=1, kernel_regularizer=keras.regularizers.l2(l=1e-4))(instance_pred)
 
 CAP and RAP can be combined, of course, by applying both a kernel constraint
 and a kernel regularizer.
diff --git a/docs/index.rst b/docs/index.rst
index 1dfdc5b..88e0567 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -10,9 +10,8 @@ AutoPool
 ========
 
 AutoPool is an adaptive pooling function with a learnable parameter that allows
-it to smoothly interpolate between min-, mean-, softmax- and max-pooling. It is
-implemented as a `Keras <https://keras.io/>`_ layer for easy integration in
-deep learning architectures.
+it to smoothly interpolate between min-, mean-, softmax- and max-pooling.
+It is implemented as `Keras <https://keras.io/>`_ or `Tensorflow-Keras <https://www.tensorflow.org/api_docs/python/tf/keras>`_ layers for easy integration in deep learning architectures.
 
 AutoPool was originally developed for multiple instance learning from weakly
 labeled time-series data as detailed in the
diff --git a/docs/installation.rst b/docs/installation.rst
index a5352d9..3e340b1 100644
--- a/docs/installation.rst
+++ b/docs/installation.rst
@@ -9,17 +9,12 @@ The simplest way to install *autopool* is through the Python Package Index
 (PyPI). This will ensure that all required dependencies are fulfilled.
 This can be achieved by executing the following command::
 
-    pip install autopool
+    python -m pip install autopool[keras]
 
-or::
+which installs autopool for the Keras backend.
+If you would rather use tensorflow, you can do the following::
 
-    sudo pip install autopool
-
-to install system-wide, or::
-
-    pip install -u autopool
-
-to install just for your own user.
+    python -m pip install autopool[tf]
 
 Source
 ~~~~~~
@@ -37,8 +32,8 @@ install with `pip install -e` to link to your actively developed source tree::
 
     tar xzf autopool-VERSION.tar.gz
     cd autopool-VERSION/
-    pip install -e .
+    python -m pip install -e .
 
 Alternately, the latest development version can be installed via pip::
 
-    pip install git+https://github.com/marl/autopool
+    python -m pip install git+https://github.com/marl/autopool
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..af4c9af
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,6 @@
+[build-system]
+requires = [
+    "setuptools >= 48",
+    "wheel >= 0.29.0",
+]
+build-backend = 'setuptools.build_meta'
diff --git a/setup.cfg b/setup.cfg
new file mode 100644
index 0000000..83f556e
--- /dev/null
+++ b/setup.cfg
@@ -0,0 +1,63 @@
+[pycodestyle]
+max-line-length=119
+
+[tool:pytest]
+addopts = --cov-report term-missing --cov autopool --cov-report=xml
+xfail_strict = true
+filterwarnings =
+    ignore:Using a non-tuple sequence:FutureWarning:scipy.*
+
+
+[coverage:report]
+omit =
+    */python?.?/*
+
+
+[metadata]
+name = autopool
+version = attr: autopool.version.version
+description = Adaptive pooling operators for multiple instance learning
+long_description = file: README.md
+long_description_content_type = text/markdown; charset=UTF-8
+url = https://github.com/marl/autopool
+author = MARL
+author_email = brian.mcfee@nyu.edu
+license = ISC
+license_file = LICENSE.md
+license_file_content_type = text/markdown; charset=UTF-8
+project_urls = 
+    Source = https://github.com/marl/autopool
+    Download = https://github.com/marl/autopool/releases
+classifiers =
+    License :: OSI Approved :: MIT License
+    Programming Language :: Python
+    Development Status :: 3 - Alpha
+    Intended Audience :: Developers
+    Topic :: Software Development
+    Programming Language :: Python :: 3
+    Programming Language :: Python :: 3.6
+    Programming Language :: Python :: 3.7
+    Programming Language :: Python :: 3.8
+    Programming Language :: Python :: 3.9
+
+[options]
+packages = find:
+python_requires = >= 3.6
+
+[options.extras_require]
+docs  =
+    numpydoc
+    sphinx
+    sphinx-gallery
+tests = 
+    pytest
+    pytest-cov
+    keras >= 2.2
+    tensorflow >= 2.0
+keras = 
+    keras >= 2.2
+tf = 
+    tensorflow >= 2.0
+all =
+    keras >= 2.2
+    tensorflow >= 2.0
diff --git a/setup.py b/setup.py
index 86091cb..6b40b52 100644
--- a/setup.py
+++ b/setup.py
@@ -1,41 +1,4 @@
-from setuptools import setup, find_packages
+from setuptools import setup
 
-import imp
-
-version = imp.load_source('autopool.version', 'autopool/version.py')
-description='Adaptive pooling operators for multiple instance learning'
-
-with open('README.md') as file:
-    long_description = file.read()
-
-setup(
-    name='autopool',
-    version=version.version,
-    description=description,
-    author='MARL',
-    author_email='brian.mcfee@nyu.edu',
-    url='http://github.com/marl/autopool',
-    download_url='http://github.com/marl/autopool/releases',
-    packages=find_packages(),
-    long_description=long_description,
-    long_description_content_type='text/markdown',
-    classifiers=[
-        "License :: OSI Approved :: ISC License (ISCL)",
-        "Programming Language :: Python",
-        "Development Status :: 3 - Alpha",
-        "Intended Audience :: Developers",
-        "Topic :: Multimedia :: Sound/Audio :: Analysis",
-        "Programming Language :: Python :: 2",
-        "Programming Language :: Python :: 2.7",
-        "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.4",
-        "Programming Language :: Python :: 3.5",
-        "Programming Language :: Python :: 3.6",
-    ],
-    keywords='deep learning, multiple instance learning, machine learning',
-    license='MIT',
-    install_requires=[
-        'keras>=2.0',
-    ],
-    extras_require={ 'docs': ['numpydoc', 'sphinx-gallery'] }
-)
+if __name__ == '__main__':
+    setup()


Test patch:
diff --git a/tests/test_keras.py b/tests/test_keras.py
new file mode 100644
index 0000000..7069969
--- /dev/null
+++ b/tests/test_keras.py
@@ -0,0 +1,58 @@
+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+
+import pytest
+import numpy as np
+import autopool.keras
+
+
+@pytest.mark.parametrize('init', ['zeros', 'ones'])
+def test_autopool_keras(init):
+
+    # Create a random array
+    x = np.random.randn(1, 3, 10)**2
+
+    # pool over the trailing dimension
+    ap_layer = autopool.keras.AutoPool1D(kernel_initializer=init, axis=-1)
+
+    x_ap = ap_layer(x)
+
+    assert x_ap.shape == (1, 3)
+
+    # Compute by force with numpy
+    if init == 'zeros':
+        # For zero-init, the default state will be a normal average
+        scaled = np.zeros_like(x)
+    elif init == 'ones':
+        # For ones-init, we should have the same behavior as softmax pooling
+        scaled = x
+
+    max_val = np.max(scaled, axis=-1, keepdims=True)
+    softmax = np.exp(scaled - max_val)
+    weights = softmax / np.sum(softmax, axis=-1, keepdims=True)
+
+    x_np_ap = np.sum(x * weights, axis=-1, keepdims=False)
+
+    assert np.allclose(x_np_ap, x_ap), (x, x_ap, x_np_ap)
+
+
+def test_softmaxpool_keras():
+    # Create a random array
+    x = np.random.randn(1, 3, 10)**2
+
+    # pool over the trailing dimension
+    ap_layer = autopool.keras.SoftMaxPool1D(axis=-1)
+
+    x_ap = ap_layer(x)
+
+    assert x_ap.shape == (1, 3)
+
+    # Compute by force with numpy
+    max_val = np.max(x, axis=-1, keepdims=True)
+    softmax = np.exp(x - max_val)
+    weights = softmax / np.sum(softmax, axis=-1, keepdims=True)
+
+    x_np_ap = np.sum(x * weights, axis=-1, keepdims=False)
+
+    assert np.allclose(x_np_ap, x_ap), (x, x_ap, x_np_ap)
+
diff --git a/tests/test_tf.py b/tests/test_tf.py
new file mode 100644
index 0000000..617adf1
--- /dev/null
+++ b/tests/test_tf.py
@@ -0,0 +1,58 @@
+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+
+import pytest
+import numpy as np
+import autopool.tf
+
+
+@pytest.mark.parametrize('init', ['zeros', 'ones'])
+def test_autopool_tf(init):
+
+    # Create a random array
+    x = np.random.randn(1, 3, 10)**2
+
+    # pool over the trailing dimension
+    ap_layer = autopool.tf.AutoPool1D(kernel_initializer=init, axis=-1)
+
+    x_ap = ap_layer(x)
+
+    assert x_ap.shape == (1, 3)
+
+    # Compute by force with numpy
+    if init == 'zeros':
+        # For zero-init, the default state will be a normal average
+        scaled = np.zeros_like(x)
+    elif init == 'ones':
+        # For ones-init, we should have the same behavior as softmax pooling
+        scaled = x
+
+    max_val = np.max(scaled, axis=-1, keepdims=True)
+    softmax = np.exp(scaled - max_val)
+    weights = softmax / np.sum(softmax, axis=-1, keepdims=True)
+
+    x_np_ap = np.sum(x * weights, axis=-1, keepdims=False)
+
+    assert np.allclose(x_np_ap, x_ap), (x, x_ap, x_np_ap)
+
+
+def test_softmaxpool_tf():
+    # Create a random array
+    x = np.random.randn(1, 3, 10)**2
+
+    # pool over the trailing dimension
+    ap_layer = autopool.tf.SoftMaxPool1D(axis=-1)
+
+    x_ap = ap_layer(x)
+
+    assert x_ap.shape == (1, 3)
+
+    # Compute by force with numpy
+    max_val = np.max(x, axis=-1, keepdims=True)
+    softmax = np.exp(x - max_val)
+    weights = softmax / np.sum(softmax, axis=-1, keepdims=True)
+
+    x_np_ap = np.sum(x * weights, axis=-1, keepdims=False)
+
+    assert np.allclose(x_np_ap, x_ap), (x, x_ap, x_np_ap)
+