feat: validate the complex channel axis; fix the documentation build

.github/workflows/test_python.yml

@@ -37,6 +37,25 @@ jobs:
           python -m pip install dist/*.whl
           python -c "import complexnn"
 
+  # Build the documentation the way Read the Docs does: install the package
+  # and the docs requirements. -W turns every Sphinx warning into an error.
+  docs:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v7
+        with:
+          fetch-depth: 0
+      - uses: astral-sh/setup-uv@v10.3.0
+        with:
+          python-version: "3.12"
+      - name: Build the documentation
+        env:
+          TF_CPP_MIN_LOG_LEVEL: "3"
+        run: >-
+          uv run --no-project
+          --with-requirements docs/requirements.txt --with-editable .
+          sphinx-build -W --keep-going -b html docs/source docs/build/html
+
   test:
     name: ${{ matrix.name }}
     runs-on: ubuntu-latest

.gitignore

@@ -48,4 +48,8 @@ venv
 
 dist/*
 
-__pycache__/*
+__pycache__/*
+
+# Sphinx output and the API pages sphinxcontrib-apidoc generates
+docs/build/
+docs/source/api/

.readthedocs.yml

@@ -7,9 +7,9 @@ version: 2
 
 # Set the version of Python and other tools you might need
 build:
-  os: ubuntu-20.04
+  os: ubuntu-24.04
   tools:
-    python: "3.10"
+    python: "3.12"
     # You can also specify other tool versions:
     # nodejs: "16"
     # rust: "1.55"
@@ -23,6 +23,9 @@ sphinx:
 formats:
    - pdf
 
+# Install the package so autodoc can import it, then the Sphinx extensions.
 python:
    install:
-   - requirements: docs/requirements.txt
+   - method: pip
+     path: .
+   - requirements: docs/requirements.txt

README.md

@@ -51,6 +51,8 @@ model.compile(optimizer=optimizers.Adam(), loss='mse')
 
 ```
 
+The layers take complex data as real tensors, with the real and imaginary parts stacked along the channel axis: the first half of the channels holds the real parts and the second half the imaginary parts. `input_shape=(28, 28, 2)` above is therefore a 28x28 image with one complex channel. The channel axis must have an even size of at least 2, and the layers raise a `ValueError` when it does not. [Complex Format of Tensors](#complex-format-of-tensors) below has the details.
+
 An example working implementation of an autoencoder can be found [here](https://github.com/JesperDramsch/Complex-CNN-Seismic/).
 
 Saving and loading
@@ -69,20 +71,23 @@ model = tf.keras.models.load_model("model.keras")
 Complex Format of Tensors
 -------------------------
 
-This library assumes that complex values are split into two real-valued parts. The real-valued and complex-valued complement, also seen [in the Docs](https://keras-complex.readthedocs.io/math.html).
+This library stores each complex value as two real values, its real part and its imaginary part. The two parts are stacked along the channel axis, all real parts first:
 
-The tensors for a 2D complex tensor of 3x3, the look like:
+- along the last axis with `data_format="channels_last"`, the default;
+- along axis 1 with `data_format="channels_first"`.
 
-```
-[[[r r r],
-  [r r r],
-  [r r r]],
-  [i,i,i],
-  [i,i,i],
-  [i,i,i]]]
+For a batch of 3x3 images with two complex channels `z0` and `z1`, the input has the shape `(batch, 3, 3, 4)`, and its channels hold `[Re z0, Re z1, Im z0, Im z1]` in this order. The parts are not interleaved.
+
+To turn a complex NumPy array into this format:
+
+```python
+import numpy as np
+
+z = np.random.randn(8, 3, 3, 2) + 1j * np.random.randn(8, 3, 3, 2)  # two complex channels
+x = np.concatenate([z.real, z.imag], axis=-1)  # shape (8, 3, 3, 4)
 ```
 
-So multiple samples should then be arranged into `[r,r,r,i,i,i]`, which is also documented [in the Docs](https://keras-complex.readthedocs.io/math.html#implementation).
+The complex layers return their outputs in the same format: `ComplexConv2D(filters=f)` returns `2 * f` channels, the real parts first. `complexnn.GetReal`, `complexnn.GetImag` and `complexnn.GetAbs` split such a tensor into its real parts, imaginary parts and magnitudes. The [data layout page in the docs](https://keras-complex.readthedocs.io/layout.html) lists the axis each layer uses, and the [math page](https://keras-complex.readthedocs.io/math.html#implementation) shows how the complex convolution is computed.
 
 Citation
 --------

complexnn/bn.py

@@ -16,6 +16,8 @@
 
 from .utils import _ignore_input_shape_warning
 
+from .utils import _check_complex_axis
+
 
 def sqrt_init(shape, dtype=None):
     value = (1 / np.sqrt(2)) * K.ones(shape)
@@ -268,12 +270,16 @@ class ComplexBatchNormalization(Layer):
     close to the null vector, the 2 by 2 covariance matrix of a complex unit close to identity
     and the 2 by 2 relation matrix, also called pseudo-covariance, close to the
     null matrix.
-    # Arguments
+
+    Arguments:
         axis: Integer, the axis that should be normalized
-            (typically the features axis).
+            (typically the features axis). It holds the real parts
+            followed by the imaginary parts, so its size must be even.
             For instance, after a `Conv2D` layer with
             `data_format="channels_first"`,
-            set `axis=2` in `ComplexBatchNormalization`.
+            set `axis=1` in `ComplexBatchNormalization`.
+            Only `axis=-1` and, for inputs of rank 2, 4 or 5,
+            `axis=1` are supported.
         momentum: Momentum for the moving statistics related to the real and
             imaginary parts.
         epsilon: Small float added to each of the variances related to the
@@ -295,13 +301,16 @@ class ComplexBatchNormalization(Layer):
         gamma_regularizer: Optional regularizer for the gamma weights.
         beta_constraint: Optional constraint for the beta weights.
         gamma_constraint: Optional constraint for the gamma weights.
-    # Input shape
+
+    Input shape:
         Arbitrary. Use the keyword argument `input_shape`
         (tuple of integers, does not include the samples axis)
         when using this layer as the first layer in a model.
-    # Output shape
+
+    Output shape:
         Same shape as input.
-    # References
+
+    References:
         - [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://arxiv.org/abs/1502.03167)
     """
 
@@ -359,6 +368,13 @@ def build(self, input_shape):
                 "input tensor should have a defined dimension "
                 "but the layer received an input with shape " + str(input_shape) + "."
             )
+        _check_complex_axis(
+            self,
+            input_shape,
+            self.axis,
+            hint="Use `axis=-1` when the parts are on the last axis, or `axis=1` "
+            "when they are on axis 1 of an input of rank 2, 4 or 5.",
+        )
         self.input_spec = InputSpec(ndim=len(input_shape), axes={self.axis: dim})
 
         param_shape = (input_shape[self.axis] // 2,)

complexnn/conv.py

@@ -15,7 +15,7 @@
 from .bn import ComplexBN as complex_normalization
 from .bn import sqrt_init
 from .init import ComplexIndependentFilters, ComplexInit
-from .utils import _ignore_input_shape_warning, get_initializer
+from .utils import _check_complex_axis, _ignore_input_shape_warning, get_initializer
 
 # WeightNorm_Conv still subclasses the Keras 2 convolution base from TensorFlow's
 # private copy of Keras. The class is broken on Keras 2 and Keras 3 and needs a
@@ -356,6 +356,13 @@ def build(self, input_shape):
             raise ValueError(
                 "The channel dimension of the inputs should be defined. Found `None`."
             )
+        _check_complex_axis(
+            self,
+            input_shape,
+            channel_axis,
+            hint=f"The channel axis is axis {channel_axis} because data_format "
+            f"is '{self.data_format}'.",
+        )
         # Divide by 2 for real and complex input.
         input_dim = input_shape[channel_axis] // 2
         if False and self.transposed:
@@ -658,9 +665,10 @@ class ComplexConv1D(ComplexConv):
     When using this layer as the first layer in a model,
     provide an `input_shape` argument
     (tuple of integers or `None`, e.g.
-    `(10, 128)` for sequences of 10 vectors of 128-dimensional vectors,
-    or `(None, 128)` for variable-length sequences of 128-dimensional vectors.
-    # Arguments
+    `(10, 128)` for sequences of 10 vectors of 64 complex values each,
+    or `(None, 128)` for variable-length sequences of such vectors.
+
+    Arguments:
         filters: Integer, the dimensionality of the output space, i.e,
             the number of complex feature maps. It is also the effective number
             of feature maps for each of the real and imaginary parts.
@@ -718,10 +726,15 @@ class ComplexConv1D(ComplexConv):
         spectral_parametrization: Whether or not to use a spectral
             parametrization of the parameters.
         transposed: Boolean, whether or not to use transposed convolution
-    # Input shape
-        3D tensor with shape: `(batch_size, steps, input_dim)`
-    # Output shape
-        3D tensor with shape: `(batch_size, new_steps, 2 x filters)`
+
+    Input shape:
+        3D tensor with shape: `(batch_size, steps, input_dim)`. The first
+        half of `input_dim` holds the real parts and the second half the
+        imaginary parts, so `input_dim` must be even.
+
+    Output shape:
+        3D tensor with shape: `(batch_size, new_steps, 2 x filters)`,
+        with the real parts in the first `filters` channels.
         `steps` value might have changed due to padding or strides.
     """
 
@@ -790,9 +803,10 @@ class ComplexConv2D(ComplexConv):
     When using this layer as the first layer in a model,
     provide the keyword argument `input_shape`
     (tuple of integers, does not include the sample axis),
-    e.g. `input_shape=(128, 128, 3)` for 128x128 RGB pictures
-    in `data_format="channels_last"`.
-    # Arguments
+    e.g. `input_shape=(128, 128, 6)` for 128x128 images with 3 complex
+    channels in `data_format="channels_last"`.
+
+    Arguments:
         filters: Integer, the dimensionality of the complex output space
             (i.e, the number complex feature maps in the convolution). The
             total effective number of filters or feature maps is 2 x filters.
@@ -860,18 +874,23 @@ class ComplexConv2D(ComplexConv):
         spectral_parametrization: Whether or not to use a spectral
             parametrization of the parameters.
         transposed: Boolean, whether or not to use transposed convolution
-    # Input shape
+
+    Input shape:
         4D tensor with shape:
         `(samples, channels, rows, cols)` if data_format='channels_first'
         or 4D tensor with shape:
         `(samples, rows, cols, channels)` if data_format='channels_last'.
-    # Output shape
+        `channels` holds the real parts of the complex channels followed by
+        their imaginary parts, so it must be even.
+
+    Output shape:
         4D tensor with shape:
         `(samples, 2 x filters, new_rows, new_cols)` if
         data_format='channels_first' or 4D tensor with shape:
         `(samples, new_rows, new_cols, 2 x filters)` if
         data_format='channels_last'. `rows` and `cols` values might have
-        changed due to padding.
+        changed due to padding. The first `filters` channels hold the real
+        parts and the last `filters` channels the imaginary parts.
     """
 
     def __init__(
@@ -938,10 +957,10 @@ class ComplexConv3D(ComplexConv):
     When using this layer as the first layer in a model,
     provide the keyword argument `input_shape`
     (tuple of integers, does not include the sample axis),
-    e.g. `input_shape=(2, 128, 128, 128, 3)` for 128x128x128 volumes
-    with 3 channels,
-    in `data_format="channels_last"`.
-    # Arguments
+    e.g. `input_shape=(128, 128, 128, 2)` for 128x128x128 volumes
+    with one complex channel, in `data_format="channels_last"`.
+
+    Arguments:
         filters: Integer, the dimensionality of the complex output space
             (i.e, the number complex feature maps in the convolution). The
             total effective number of filters or feature maps is 2 x filters.
@@ -1008,22 +1027,27 @@ class ComplexConv3D(ComplexConv):
         spectral_parametrization: Whether or not to use a spectral
             parametrization of the parameters.
         transposed: Boolean, whether or not to use transposed convolution
-    # Input shape
+
+    Input shape:
         5D tensor with shape:
         `(samples, channels, conv_dim1, conv_dim2, conv_dim3)` if
         data_format='channels_first'
         or 5D tensor with shape:
         `(samples, conv_dim1, conv_dim2, conv_dim3, channels)` if
         data_format='channels_last'.
-    # Output shape
+        `channels` holds the real parts of the complex channels followed by
+        their imaginary parts, so it must be even.
+
+    Output shape:
         5D tensor with shape:
         `(samples, 2 x filters, new_conv_dim1, new_conv_dim2, new_conv_dim3)`
         if data_format='channels_first'
         or 5D tensor with shape:
         `(samples, new_conv_dim1, new_conv_dim2, new_conv_dim3, 2 x filters)`
         if data_format='channels_last'.
         `new_conv_dim1`, `new_conv_dim2` and `new_conv_dim3` values might have
-        changed due to padding.
+        changed due to padding. The first `filters` channels hold the real
+        parts and the last `filters` channels the imaginary parts.
     """
 
     def __init__(

complexnn/dense.py

@@ -8,7 +8,12 @@
 from tensorflow.keras.layers import InputSpec, Layer
 from tensorflow.keras.utils import register_keras_serializable
 
-from .utils import _compute_fans, _ignore_input_shape_warning, get_initializer
+from .utils import (
+    _check_complex_axis,
+    _compute_fans,
+    _ignore_input_shape_warning,
+    get_initializer,
+)
 
 
 @register_keras_serializable(package="complexnn")
@@ -24,7 +29,8 @@ class ComplexDense(Layer):
     (only applicable if `use_bias` is `True`).
     Note: if the input to the layer has a rank greater than 2, then
     AN ERROR MESSAGE IS PRINTED.
-    # Arguments
+
+    Arguments:
         units: Positive integer, dimensionality of each of the real part
             and the imaginary part. It is actualy the number of complex units.
         activation: Activation function to use
@@ -50,11 +56,16 @@ class ComplexDense(Layer):
             (see keras.constraints).
         bias_constraint: Constraint function applied to the bias vector
             (see keras.constraints).
-    # Input shape
-        a 2D input with shape `(batch_size, input_dim)`.
-    # Output shape
+
+    Input shape:
+        a 2D input with shape `(batch_size, input_dim)`. The first half of
+        `input_dim` holds the real parts and the second half the imaginary
+        parts, so `input_dim` must be even.
+
+    Output shape:
         For a 2D input with shape `(batch_size, input_dim)`,
-        the output would have shape `(batch_size, units)`.
+        the output would have shape `(batch_size, 2 x units)`,
+        with the real parts in the first `units` features.
     """
 
     def __init__(
@@ -102,7 +113,7 @@ def __init__(
 
     def build(self, input_shape):
         assert len(input_shape) == 2
-        assert input_shape[-1] % 2 == 0
+        _check_complex_axis(self, input_shape, -1)
         input_dim = input_shape[-1] // 2
         data_format = K.image_data_format()
         kernel_shape = (input_dim, self.units)

complexnn/norm.py

@@ -11,7 +11,7 @@
 
 from .bn import ComplexBN as complex_normalization
 from .bn import sanitizedInitGet, sanitizedInitSer, sqrt_init
-from .utils import _ignore_input_shape_warning
+from .utils import _check_complex_axis, _ignore_input_shape_warning
 
 
 def layernorm(x, axis, epsilon, gamma, beta):
@@ -154,6 +154,13 @@ def build(self, input_shape):
                 "input tensor should have a defined dimension "
                 "but the layer received an input with shape " + str(input_shape) + "."
             )
+        _check_complex_axis(
+            self,
+            input_shape,
+            self.axis,
+            hint="Use `axis=-1` when the parts are on the last axis, or `axis=1` "
+            "when they are on axis 1 of an input of rank 2, 4 or 5.",
+        )
         self.input_spec = InputSpec(ndim=len(input_shape), axes={self.axis: dim})
 
         gamma_shape = (input_shape[self.axis] // 2,)