From a606be472ac6bdb7f0c8571f6e6e6fe3bc6799c0 Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Mon, 12 Oct 2026 14:37:56 +0200 Subject: [PATCH 01/16] style(conv): join split error strings as ruff format does Current ruff format joins the three implicitly concatenated "channel dimension" messages onto one line. Applying it here lets the following commits on conv.py pass "ruff format --check" without carrying unrelated formatting changes. No behaviour change. --- complexnn/conv.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/complexnn/conv.py b/complexnn/conv.py index 1f6303d..759f8c2 100644 --- a/complexnn/conv.py +++ b/complexnn/conv.py @@ -333,8 +333,7 @@ def build(self, input_shape): channel_axis = -1 if input_shape[channel_axis] is None: raise ValueError( - "The channel dimension of the inputs " - "should be defined. Found `None`." + "The channel dimension of the inputs should be defined. Found `None`." ) # Divide by 2 for real and complex input. input_dim = input_shape[channel_axis] // 2 @@ -1106,8 +1105,7 @@ def build(self, input_shape): channel_axis = -1 if input_shape[channel_axis] is None: raise ValueError( - "The channel dimension of the inputs " - "should be defined. Found `None`." + "The channel dimension of the inputs should be defined. Found `None`." ) input_dim = input_shape[channel_axis] gamma_shape = (input_dim * self.filters,) @@ -1127,8 +1125,7 @@ def call(self, inputs): channel_axis = -1 if input_shape[channel_axis] is None: raise ValueError( - "The channel dimension of the inputs " - "should be defined. Found `None`." + "The channel dimension of the inputs should be defined. Found `None`." ) input_dim = input_shape[channel_axis] ker_shape = self.kernel_size + (input_dim, self.filters) From 599b807f81a6de678b7ee2e736d6a04be07ab611 Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Mon, 12 Oct 2026 14:40:58 +0200 Subject: [PATCH 02/16] fix(conv): leave data_format out of the ComplexConv1D config ComplexConv1D.__init__ passes data_format="channels_last" to ComplexConv itself, but get_config still wrote data_format, so from_config and .keras loading failed with "got multiple values for keyword argument 'data_format'" on Keras 2 and Keras 3. Drop it from the config, as is already done for rank. --- complexnn/conv.py | 2 ++ tests/test_conv.py | 12 ++++++++++++ 2 files changed, 14 insertions(+) diff --git a/complexnn/conv.py b/complexnn/conv.py index 759f8c2..cee3749 100644 --- a/complexnn/conv.py +++ b/complexnn/conv.py @@ -768,6 +768,8 @@ def __init__( def get_config(self): config = super(ComplexConv1D, self).get_config() config.pop("rank") + # __init__ always passes data_format="channels_last" to ComplexConv. + config.pop("data_format") return config diff --git a/tests/test_conv.py b/tests/test_conv.py index 4d82ecb..05499e5 100644 --- a/tests/test_conv.py +++ b/tests/test_conv.py @@ -56,3 +56,15 @@ def test_conv2Dtranspose(self): true = (None, 128, 128, 4) calc = model.output_shape self.assertEqual(true, calc) + + +class TestConvConfig(unittest.TestCase): + """Conv layer configs rebuild the same layer""" + + def test_conv1d_from_config(self): + """ComplexConv1D sets data_format itself, so its config must not repeat it""" + layer = conn.ComplexConv1D(filters=2, kernel_size=3, padding="same") + config = layer.get_config() + clone = conn.ComplexConv1D.from_config(config) + self.assertEqual(clone.get_config(), config) + self.assertEqual(clone.data_format, "channels_last") From 9a558dfaac6cf2b4f6d1f870973b68b36591c647 Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Mon, 12 Oct 2026 14:41:44 +0200 Subject: [PATCH 03/16] fix(conv): store epsilon in the ComplexConv config With normalize_weight=True, epsilon is added to the kernel variances before normalisation, so it changes the layer output. get_config left it out, and a layer rebuilt from its config fell back to the default 1e-7. --- complexnn/conv.py | 1 + tests/test_conv.py | 8 ++++++++ 2 files changed, 9 insertions(+) diff --git a/complexnn/conv.py b/complexnn/conv.py index cee3749..a51c36b 100644 --- a/complexnn/conv.py +++ b/complexnn/conv.py @@ -635,6 +635,7 @@ def get_config(self): "init_criterion": self.init_criterion, "spectral_parametrization": self.spectral_parametrization, "transposed": self.transposed, + "epsilon": self.epsilon, } base_config = super(ComplexConv, self).get_config() return dict(list(base_config.items()) + list(config.items())) diff --git a/tests/test_conv.py b/tests/test_conv.py index 05499e5..c83da31 100644 --- a/tests/test_conv.py +++ b/tests/test_conv.py @@ -68,3 +68,11 @@ def test_conv1d_from_config(self): clone = conn.ComplexConv1D.from_config(config) self.assertEqual(clone.get_config(), config) self.assertEqual(clone.data_format, "channels_last") + + def test_epsilon_in_config(self): + """epsilon changes the normalised kernel, so from_config must keep it""" + layer = conn.ComplexConv2D( + filters=2, kernel_size=3, normalize_weight=True, epsilon=1e-3 + ) + clone = conn.ComplexConv2D.from_config(layer.get_config()) + self.assertEqual(clone.epsilon, 1e-3) From 5870a35fd6835f3db221fd1b08d39b7dbb056066 Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Mon, 12 Oct 2026 14:44:46 +0200 Subject: [PATCH 04/16] fix(norm): store the sqrt_init default of ComplexLayerNorm by name The default gamma_diag_initializer is the plain function bn.sqrt_init. initializers.serialize wrote it as a "builtins.function" entry that initializers.get cannot read back, so from_config and .keras loading of ComplexLayerNorm failed on Keras 3, and the config did not round-trip on Keras 2 either. Use sanitizedInitGet and sanitizedInitSer from bn.py, as ComplexBatchNormalization and ComplexConv already do, so the function is stored as "sqrt_init" and other initialisers still go through Keras. --- complexnn/norm.py | 18 +++++++++--------- tests/test_norm.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 9 deletions(-) create mode 100644 tests/test_norm.py diff --git a/complexnn/norm.py b/complexnn/norm.py index d66e193..c9709de 100644 --- a/complexnn/norm.py +++ b/complexnn/norm.py @@ -10,7 +10,7 @@ from tensorflow.keras.layers import InputSpec, Layer from .bn import ComplexBN as complex_normalization -from .bn import sqrt_init +from .bn import sanitizedInitGet, sanitizedInitSer, sqrt_init def layernorm(x, axis, epsilon, gamma, beta): @@ -129,9 +129,11 @@ def __init__( self.axis = axis self.center = center self.scale = scale - self.beta_initializer = initializers.get(beta_initializer) - self.gamma_diag_initializer = initializers.get(gamma_diag_initializer) - self.gamma_off_initializer = initializers.get(gamma_off_initializer) + # sanitizedInitGet also accepts "sqrt_init", the name under which + # get_config stores the default gamma_diag_initializer. + self.beta_initializer = sanitizedInitGet(beta_initializer) + self.gamma_diag_initializer = sanitizedInitGet(gamma_diag_initializer) + self.gamma_off_initializer = sanitizedInitGet(gamma_off_initializer) self.beta_regularizer = regularizers.get(beta_regularizer) self.gamma_diag_regularizer = regularizers.get(gamma_diag_regularizer) self.gamma_off_regularizer = regularizers.get(gamma_off_regularizer) @@ -272,11 +274,9 @@ def get_config(self): "epsilon": self.epsilon, "center": self.center, "scale": self.scale, - "beta_initializer": initializers.serialize(self.beta_initializer), - "gamma_diag_initializer": initializers.serialize( - self.gamma_diag_initializer - ), - "gamma_off_initializer": initializers.serialize(self.gamma_off_initializer), + "beta_initializer": sanitizedInitSer(self.beta_initializer), + "gamma_diag_initializer": sanitizedInitSer(self.gamma_diag_initializer), + "gamma_off_initializer": sanitizedInitSer(self.gamma_off_initializer), "beta_regularizer": regularizers.serialize(self.beta_regularizer), "gamma_diag_regularizer": regularizers.serialize( self.gamma_diag_regularizer diff --git a/tests/test_norm.py b/tests/test_norm.py new file mode 100644 index 0000000..b240a4a --- /dev/null +++ b/tests/test_norm.py @@ -0,0 +1,28 @@ +import unittest + +import complexnn as conn +import tensorflow as tf +from complexnn.bn import sqrt_init + + +class TestComplexLayerNormConfig(unittest.TestCase): + """ComplexLayerNorm configs rebuild the same layer""" + + def test_from_config(self): + layer = conn.ComplexLayerNorm() + config = layer.get_config() + clone = conn.ComplexLayerNorm.from_config(config) + self.assertEqual(clone.get_config(), config) + + def test_sqrt_init_serialises_by_name(self): + """The default gamma initialiser is stored by name and restored as sqrt_init""" + config = conn.ComplexLayerNorm().get_config() + self.assertEqual(config["gamma_diag_initializer"], "sqrt_init") + clone = conn.ComplexLayerNorm.from_config(config) + self.assertIs(clone.gamma_diag_initializer, sqrt_init) + + def test_keras_initializer_round_trip(self): + """Keras initialisers still go through the Keras registry""" + layer = conn.ComplexLayerNorm(gamma_diag_initializer="ones") + clone = conn.ComplexLayerNorm.from_config(layer.get_config()) + self.assertIsInstance(clone.gamma_diag_initializer, tf.keras.initializers.Ones) From 5ebc6972ddf12ba00738444e021853dcdc53d5d5 Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Mon, 12 Oct 2026 14:45:38 +0200 Subject: [PATCH 05/16] fix(norm): keep initialisers and regularisers in the LayerNormalization config get_config dropped beta_init and gamma_init, so a rebuilt layer always used the defaults, and it stored each regulariser as its bare get_config() dict, which regularizers.get cannot read back. Serialise all four with the Keras serialize functions instead. --- complexnn/norm.py | 10 ++++------ tests/test_norm.py | 17 +++++++++++++++++ 2 files changed, 21 insertions(+), 6 deletions(-) diff --git a/complexnn/norm.py b/complexnn/norm.py index c9709de..d36fc16 100644 --- a/complexnn/norm.py +++ b/complexnn/norm.py @@ -95,12 +95,10 @@ def get_config(self): config = { "epsilon": self.epsilon, "axis": self.axis, - "gamma_regularizer": self.gamma_regularizer.get_config() - if self.gamma_regularizer - else None, - "beta_regularizer": self.beta_regularizer.get_config() - if self.beta_regularizer - else None, + "beta_init": initializers.serialize(self.beta_init), + "gamma_init": initializers.serialize(self.gamma_init), + "gamma_regularizer": regularizers.serialize(self.gamma_regularizer), + "beta_regularizer": regularizers.serialize(self.beta_regularizer), } base_config = super(LayerNormalization, self).get_config() return dict(list(base_config.items()) + list(config.items())) diff --git a/tests/test_norm.py b/tests/test_norm.py index b240a4a..2e15207 100644 --- a/tests/test_norm.py +++ b/tests/test_norm.py @@ -26,3 +26,20 @@ def test_keras_initializer_round_trip(self): layer = conn.ComplexLayerNorm(gamma_diag_initializer="ones") clone = conn.ComplexLayerNorm.from_config(layer.get_config()) self.assertIsInstance(clone.gamma_diag_initializer, tf.keras.initializers.Ones) + + +class TestLayerNormalizationConfig(unittest.TestCase): + """LayerNormalization configs keep their initialisers and regularisers""" + + def test_from_config(self): + layer = conn.LayerNormalization( + beta_init="ones", + gamma_init="zeros", + gamma_regularizer=tf.keras.regularizers.L2(0.01), + ) + config = layer.get_config() + clone = conn.LayerNormalization.from_config(config) + self.assertEqual(clone.get_config(), config) + self.assertIsInstance(clone.beta_init, tf.keras.initializers.Ones) + self.assertIsInstance(clone.gamma_init, tf.keras.initializers.Zeros) + self.assertAlmostEqual(clone.gamma_regularizer.get_config()["l2"], 0.01) From f2b5ef92f9144fa29a2d42ffea456ac864a78b87 Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Mon, 12 Oct 2026 14:46:36 +0200 Subject: [PATCH 06/16] feat(init): add get_config to ComplexInit ComplexInit fell back to the base Initializer.get_config, which returns an empty dict, so from_config raised a TypeError for the missing kernel_size, input_dim and weight_dim. Return the constructor arguments, as IndependentFilters and ComplexIndependentFilters already do. SqrtInit takes no arguments, so the inherited empty config is correct for it. --- complexnn/init.py | 10 ++++++++++ tests/test_init.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 40 insertions(+) create mode 100644 tests/test_init.py diff --git a/complexnn/init.py b/complexnn/init.py index 6ebd7fc..665f8a3 100644 --- a/complexnn/init.py +++ b/complexnn/init.py @@ -281,6 +281,16 @@ def __call__(self, shape, dtype=None): return weight + def get_config(self): + return { + "nb_filters": self.nb_filters, + "kernel_size": self.kernel_size, + "input_dim": self.input_dim, + "weight_dim": self.weight_dim, + "criterion": self.criterion, + "seed": self.seed, + } + class SqrtInit(Initializer): def __call__(self, shape, dtype=None): diff --git a/tests/test_init.py b/tests/test_init.py new file mode 100644 index 0000000..0695ab5 --- /dev/null +++ b/tests/test_init.py @@ -0,0 +1,30 @@ +import unittest + +import complexnn as conn +import numpy as np + + +def make_initializers(): + """One of each complex initialiser, each filling a (4, 6) dense kernel""" + return [ + conn.ComplexInit( + kernel_size=(4, 6), input_dim=4, weight_dim=2, criterion="glorot", seed=3 + ), + conn.IndependentFilters(kernel_size=(4, 6), input_dim=4, weight_dim=2, seed=3), + conn.ComplexIndependentFilters( + kernel_size=(4, 3), input_dim=4, weight_dim=2, seed=3 + ), + ] + + +class TestInitializerConfig(unittest.TestCase): + """Initialiser configs rebuild an initialiser that gives the same values""" + + shape = (4, 6) + + def test_config_round_trip(self): + for init in make_initializers(): + with self.subTest(type(init).__name__): + clone = type(init).from_config(init.get_config()) + self.assertEqual(clone.get_config(), init.get_config()) + np.testing.assert_array_equal(clone(self.shape), init(self.shape)) From f1b1a02cf6f9f50dc4bb0d6cebb22041dfd1a125 Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Mon, 12 Oct 2026 14:48:11 +0200 Subject: [PATCH 07/16] feat: give ComplexBatchNormalization and ComplexLayerNorm an output shape Keras 3 raises NotImplementedError from compute_output_shape on layers that do not define it, and infers output specs by tracing call() instead. Both layers keep the input shape, so return it unchanged. Keras 2 already inferred the same shape by tracing, so its results do not change. --- complexnn/bn.py | 5 +++++ complexnn/norm.py | 5 +++++ tests/test_bn.py | 10 ++++++++++ tests/test_norm.py | 10 ++++++++++ 4 files changed, 30 insertions(+) diff --git a/complexnn/bn.py b/complexnn/bn.py index 95f4d3f..34c91f8 100644 --- a/complexnn/bn.py +++ b/complexnn/bn.py @@ -529,6 +529,11 @@ def update_moving_average(variable, value): axis=self.axis, ) + def compute_output_shape(self, input_shape): + # Normalisation keeps the shape. Defining it lets Keras 3 skip + # tracing call() to infer the output shape. + return input_shape + def get_config(self): config = { "axis": self.axis, diff --git a/complexnn/norm.py b/complexnn/norm.py index d36fc16..3a4b924 100644 --- a/complexnn/norm.py +++ b/complexnn/norm.py @@ -266,6 +266,11 @@ def call(self, inputs): axis=self.axis, ) + def compute_output_shape(self, input_shape): + # Normalisation keeps the shape. Defining it lets Keras 3 skip + # tracing call() to infer the output shape. + return input_shape + def get_config(self): config = { "axis": self.axis, diff --git a/tests/test_bn.py b/tests/test_bn.py index 834f529..67acb9f 100644 --- a/tests/test_bn.py +++ b/tests/test_bn.py @@ -48,3 +48,13 @@ def test_fit_updates_moving_mean(self): self.x, np.zeros_like(self.x), batch_size=32, epochs=1, verbose=0 ) self.assertFalse(np.allclose(self.bn.moving_mean.numpy(), before)) + + +class TestComplexBatchNormalizationShape(unittest.TestCase): + """ComplexBatchNormalization keeps the input shape""" + + def test_compute_output_shape(self): + layer = conn.ComplexBN() + for shape in [(None, 8), (None, 4, 4, 8)]: + with self.subTest(shape=shape): + self.assertEqual(tuple(layer.compute_output_shape(shape)), shape) diff --git a/tests/test_norm.py b/tests/test_norm.py index 2e15207..41544fa 100644 --- a/tests/test_norm.py +++ b/tests/test_norm.py @@ -43,3 +43,13 @@ def test_from_config(self): self.assertIsInstance(clone.beta_init, tf.keras.initializers.Ones) self.assertIsInstance(clone.gamma_init, tf.keras.initializers.Zeros) self.assertAlmostEqual(clone.gamma_regularizer.get_config()["l2"], 0.01) + + +class TestComplexLayerNormShape(unittest.TestCase): + """ComplexLayerNorm keeps the input shape""" + + def test_compute_output_shape(self): + layer = conn.ComplexLayerNorm() + for shape in [(None, 8), (None, 4, 4, 8)]: + with self.subTest(shape=shape): + self.assertEqual(tuple(layer.compute_output_shape(shape)), shape) From a84884a218feb6c777f5a7a67005e6afa16a4c84 Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Mon, 12 Oct 2026 16:44:12 +0200 Subject: [PATCH 08/16] test: rebuild each working layer from its config Add a get_config/from_config round trip for ComplexDense, ComplexConv1D/2D/3D (including normalize_weight and transposed), ComplexBatchNormalization, ComplexLayerNorm, LayerNormalization and GetReal/GetImag/GetAbs. Each rebuilt layer must report the same config and, given the original weights, compute the same output, so a config that drops an argument affecting the output is caught. The weights are moved off their initial values first, so the outputs only match if the weights were really copied. --- tests/test_saving.py | 60 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 60 insertions(+) create mode 100644 tests/test_saving.py diff --git a/tests/test_saving.py b/tests/test_saving.py new file mode 100644 index 0000000..8f60993 --- /dev/null +++ b/tests/test_saving.py @@ -0,0 +1,60 @@ +import unittest + +import complexnn as conn +import numpy as np +import tensorflow as tf + +# Every working layer, with an input shape it accepts. Each entry builds a new +# layer, so tests never share an instance. +LAYERS = { + "ComplexDense": ((8,), lambda: conn.ComplexDense(4)), + "ComplexConv1D": ((10, 4), lambda: conn.ComplexConv1D(2, 3, padding="same")), + "ComplexConv2D": ((6, 6, 4), lambda: conn.ComplexConv2D(2, 3, padding="same")), + "ComplexConv2D normalize_weight": ( + (6, 6, 4), + lambda: conn.ComplexConv2D(2, 3, padding="same", normalize_weight=True), + ), + "ComplexConv2D transposed": ( + (6, 6, 4), + lambda: conn.ComplexConv2D(1, 3, strides=2, padding="same", transposed=True), + ), + "ComplexConv3D": ((4, 4, 4, 4), lambda: conn.ComplexConv3D(2, 3, padding="same")), + "ComplexBN": ((4, 4, 8), lambda: conn.ComplexBN()), + "ComplexLayerNorm": ((8,), lambda: conn.ComplexLayerNorm()), + "LayerNormalization": ((8,), lambda: conn.LayerNormalization()), + "GetReal": ((8,), lambda: conn.GetReal()), + "GetImag": ((4, 4, 8), lambda: conn.GetImag()), + "GetAbs": ((8,), lambda: conn.GetAbs()), +} + + +def build_model(shape, make_layer, rng): + """A Sequential model around one layer, with every weight moved off its initial value""" + model = tf.keras.Sequential([tf.keras.Input(shape), make_layer()]) + # With moved weights, matching outputs show that the weights were restored + # and not just initialised again. The noise is small enough to keep the + # normalisation variances positive. + model.set_weights( + [w + rng.uniform(-0.1, 0.1, w.shape) for w in model.get_weights()] + ) + return model + + +class TestLayerConfig(unittest.TestCase): + """Each working layer rebuilds from its config and computes the same output""" + + def test_from_config(self): + rng = np.random.default_rng(0) + for name, (shape, make_layer) in LAYERS.items(): + with self.subTest(name): + layer = build_model(shape, make_layer, rng).layers[-1] + config = layer.get_config() + clone = type(layer).from_config(config) + self.assertEqual(clone.get_config(), config) + + x = rng.standard_normal((3, *shape)).astype("float32") + clone(x) # builds the clone, so it can take the weights + clone.set_weights(layer.get_weights()) + np.testing.assert_allclose( + np.asarray(clone(x)), np.asarray(layer(x)), rtol=1e-6, atol=1e-6 + ) From dd8c7ec6461042b2fdf7d2fc77b22d08b714459d Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Mon, 12 Oct 2026 16:51:13 +0200 Subject: [PATCH 09/16] feat: register layers and initialisers for Keras serialisation Models that contained complexnn layers could only be loaded from a .keras file by passing every class in custom_objects. Without it, Keras 3 could not locate the classes, and Keras 2 and tf-keras deserialised them as plain strings. norm.LayerNormalization also clashed with the built-in Keras LayerNormalization of the same name on Keras 2. Register every exported class with tensorflow.keras.utils.register_keras_serializable(package="complexnn"), which exists under that path in Keras 3, Keras 2 and tf-keras. Each class is now stored as "complexnn>ClassName", so load_model finds it once complexnn is imported. The layers and initialisers that are known not to work (WeightNorm_Conv, SpectralPooling1D/2D, SqrtInit) are registered too, so they need no change here once they are fixed. The tests check that every exported class is registered, that a .keras save and load without custom_objects gives the same predictions for each working layer, and that a built-in Dense keeps a complexnn initialiser. They silence one numpy 2 DeprecationWarning that Keras 3 raises while saving any model on TensorFlow 2.21, built-in layers included. --- complexnn/bn.py | 2 ++ complexnn/conv.py | 6 ++++ complexnn/dense.py | 2 ++ complexnn/init.py | 5 +++ complexnn/norm.py | 3 ++ complexnn/pool.py | 3 ++ complexnn/utils.py | 4 +++ tests/test_saving.py | 81 ++++++++++++++++++++++++++++++++++++++++++++ 8 files changed, 106 insertions(+) diff --git a/complexnn/bn.py b/complexnn/bn.py index 34c91f8..0ad5195 100644 --- a/complexnn/bn.py +++ b/complexnn/bn.py @@ -10,6 +10,7 @@ import tensorflow.keras.backend as K from tensorflow.keras import constraints, initializers, regularizers from tensorflow.keras.layers import InputSpec, Layer +from tensorflow.keras.utils import register_keras_serializable def sqrt_init(shape, dtype=None): @@ -254,6 +255,7 @@ def ComplexBN( return input_centred +@register_keras_serializable(package="complexnn") class ComplexBatchNormalization(Layer): """Complex version of the real domain Batch normalization layer (Ioffe and Szegedy, 2014). diff --git a/complexnn/conv.py b/complexnn/conv.py index a51c36b..923f265 100644 --- a/complexnn/conv.py +++ b/complexnn/conv.py @@ -9,6 +9,7 @@ InputSpec, Layer, ) +from tensorflow.keras.utils import register_keras_serializable from tensorflow.python.keras.layers.convolutional import Conv from tensorflow.python.keras.utils import conv_utils @@ -179,6 +180,7 @@ def sanitizedInitSer(init): return initializers.serialize(init) +@register_keras_serializable(package="complexnn") class ComplexConv(Layer): """Abstract nD complex convolution layer. @@ -641,6 +643,7 @@ def get_config(self): return dict(list(base_config.items()) + list(config.items())) +@register_keras_serializable(package="complexnn") class ComplexConv1D(ComplexConv): """1D complex convolution layer. This layer creates a complex convolution kernel that is convolved @@ -774,6 +777,7 @@ def get_config(self): return config +@register_keras_serializable(package="complexnn") class ComplexConv2D(ComplexConv): """2D Complex convolution layer (e.g. spatial convolution over images). This layer creates a complex convolution kernel that is convolved @@ -921,6 +925,7 @@ def get_config(self): return config +@register_keras_serializable(package="complexnn") class ComplexConv3D(ComplexConv): """3D convolution layer (e.g. spatial convolution over volumes). This layer creates a complex convolution kernel that is convolved @@ -1072,6 +1077,7 @@ def get_config(self): return config +@register_keras_serializable(package="complexnn") class WeightNorm_Conv(Conv): """WeightNorm_Conv""" diff --git a/complexnn/dense.py b/complexnn/dense.py index 274ffd5..003ef00 100644 --- a/complexnn/dense.py +++ b/complexnn/dense.py @@ -6,10 +6,12 @@ from tensorflow.keras import activations, constraints, initializers, regularizers from tensorflow.keras import backend as K from tensorflow.keras.layers import InputSpec, Layer +from tensorflow.keras.utils import register_keras_serializable from .utils import _compute_fans +@register_keras_serializable(package="complexnn") class ComplexDense(Layer): """Regular complex densely-connected NN layer. `Dense` implements the operation: diff --git a/complexnn/init.py b/complexnn/init.py index 665f8a3..a939c70 100644 --- a/complexnn/init.py +++ b/complexnn/init.py @@ -5,6 +5,7 @@ import tensorflow.keras.backend as K from numpy.random import RandomState from tensorflow.keras.initializers import Initializer +from tensorflow.keras.utils import register_keras_serializable from tensorflow.python.keras.utils.generic_utils import ( deserialize_keras_object, serialize_keras_object, @@ -13,6 +14,7 @@ from .utils import _compute_fans +@register_keras_serializable(package="complexnn") class IndependentFilters(Initializer): # This initialization constructs real-valued kernels # that are independent as much as possible from each other @@ -107,6 +109,7 @@ def get_config(self): } +@register_keras_serializable(package="complexnn") class ComplexIndependentFilters(Initializer): # This initialization constructs complex-valued kernels # that are independent as much as possible from each other @@ -217,6 +220,7 @@ def get_config(self): } +@register_keras_serializable(package="complexnn") class ComplexInit(Initializer): # The standard complex initialization using # either the He or the Glorot criterion. @@ -292,6 +296,7 @@ def get_config(self): } +@register_keras_serializable(package="complexnn") class SqrtInit(Initializer): def __call__(self, shape, dtype=None): return K.constant(1 / K.sqrt(2), shape=shape, dtype=dtype) diff --git a/complexnn/norm.py b/complexnn/norm.py index 3a4b924..a1711f6 100644 --- a/complexnn/norm.py +++ b/complexnn/norm.py @@ -8,6 +8,7 @@ import tensorflow.keras.backend as K from tensorflow.keras import constraints, initializers, regularizers from tensorflow.keras.layers import InputSpec, Layer +from tensorflow.keras.utils import register_keras_serializable from .bn import ComplexBN as complex_normalization from .bn import sanitizedInitGet, sanitizedInitSer, sqrt_init @@ -45,6 +46,7 @@ def layernorm(x, axis, epsilon, gamma, beta): return x_LN +@register_keras_serializable(package="complexnn") class LayerNormalization(Layer): def __init__( self, @@ -104,6 +106,7 @@ def get_config(self): return dict(list(base_config.items()) + list(config.items())) +@register_keras_serializable(package="complexnn") class ComplexLayerNorm(Layer): def __init__( self, diff --git a/complexnn/pool.py b/complexnn/pool.py index 3eac79f..4c7c670 100644 --- a/complexnn/pool.py +++ b/complexnn/pool.py @@ -6,12 +6,14 @@ import tensorflow.keras.backend as KB import tensorflow.keras.layers as KL import tensorflow.keras.optimizers as KO +from tensorflow.keras.utils import register_keras_serializable # # Spectral Pooling Layer # +@register_keras_serializable(package="complexnn") class SpectralPooling1D(KL.Layer): def __init__(self, topf=(0,), **kwargs): super(SpectralPooling1D, self).__init__() @@ -56,6 +58,7 @@ def call(self, x, mask=None): return x +@register_keras_serializable(package="complexnn") class SpectralPooling2D(KL.Layer): def __init__(self, **kwargs): super(SpectralPooling2D, self).__init__() diff --git a/complexnn/utils.py b/complexnn/utils.py index 2ce9c90..a7ffad3 100644 --- a/complexnn/utils.py +++ b/complexnn/utils.py @@ -4,6 +4,7 @@ import numpy as np import tensorflow.keras.backend as K from tensorflow.keras.layers import Lambda, Layer +from tensorflow.keras.utils import register_keras_serializable # # GetReal/GetImag Lambda layer Implementation @@ -108,6 +109,7 @@ def _compute_fans(shape, data_format="channels_last"): return fan_in, fan_out +@register_keras_serializable(package="complexnn") class GetReal(Layer): def call(self, inputs): return get_realpart(inputs) @@ -116,6 +118,7 @@ def compute_output_shape(self, input_shape): return getpart_output_shape(input_shape) +@register_keras_serializable(package="complexnn") class GetImag(Layer): def call(self, inputs): return get_imagpart(inputs) @@ -124,6 +127,7 @@ def compute_output_shape(self, input_shape): return getpart_output_shape(input_shape) +@register_keras_serializable(package="complexnn") class GetAbs(Layer): def call(self, inputs): return get_abs(inputs) diff --git a/tests/test_saving.py b/tests/test_saving.py index 8f60993..7e9b84d 100644 --- a/tests/test_saving.py +++ b/tests/test_saving.py @@ -1,8 +1,13 @@ +import inspect +import os +import tempfile import unittest +import warnings import complexnn as conn import numpy as np import tensorflow as tf +from tensorflow.keras.utils import get_registered_name, get_registered_object # Every working layer, with an input shape it accepts. Each entry builds a new # layer, so tests never share an instance. @@ -58,3 +63,79 @@ def test_from_config(self): np.testing.assert_allclose( np.asarray(clone(x)), np.asarray(layer(x)), rtol=1e-6, atol=1e-6 ) + + +class TestRegistration(unittest.TestCase): + """Every class that complexnn exports is registered for serialisation""" + + def test_public_classes_are_registered(self): + classes = [ + obj + for obj in vars(conn).values() + if inspect.isclass(obj) and obj.__module__.startswith("complexnn.") + ] + self.assertGreater(len(classes), 0) + for cls in classes: + with self.subTest(cls.__name__): + registered_name = "complexnn>" + cls.__name__ + self.assertEqual(get_registered_name(cls), registered_name) + self.assertIs(get_registered_object(registered_name), cls) + + +class TestKerasFileRoundTrip(unittest.TestCase): + """Models with complex layers save to .keras and load without custom_objects""" + + def setUp(self): + tmp = tempfile.TemporaryDirectory() + self.addCleanup(tmp.cleanup) + self.tmp = tmp.name + + def save_and_load(self, model, name): + path = os.path.join(self.tmp, name + ".keras") + with warnings.catch_warnings(): + # TensorFlow variables implement __array__ without the copy keyword + # that numpy 2 passes, so Keras 3 warns while saving any model, + # built-in layers included. The warning is not about complexnn. + warnings.filterwarnings( + "ignore", + message="__array__ implementation doesn't accept a copy keyword", + category=DeprecationWarning, + ) + model.save(path) + return tf.keras.models.load_model(path) + + def test_layers(self): + rng = np.random.default_rng(1) + for i, (name, (shape, make_layer)) in enumerate(LAYERS.items()): + with self.subTest(name): + model = build_model(shape, make_layer, rng) + loaded = self.save_and_load(model, f"layer_{i}") + self.assertIs(type(loaded.layers[-1]), type(model.layers[-1])) + + x = rng.standard_normal((3, *shape)).astype("float32") + np.testing.assert_allclose( + loaded.predict(x, verbose=0), + model.predict(x, verbose=0), + rtol=1e-6, + atol=1e-6, + ) + + def test_initializers(self): + """A built-in layer keeps a complex initialiser through saving and loading""" + initializers = [ + conn.ComplexInit(kernel_size=(4, 6), input_dim=4, weight_dim=2, seed=3), + conn.IndependentFilters( + kernel_size=(4, 6), input_dim=4, weight_dim=2, seed=3 + ), + conn.ComplexIndependentFilters( + kernel_size=(4, 3), input_dim=4, weight_dim=2, seed=3 + ), + ] + for i, init in enumerate(initializers): + with self.subTest(type(init).__name__): + dense = tf.keras.layers.Dense(6, kernel_initializer=init) + model = tf.keras.Sequential([tf.keras.Input((4,)), dense]) + loaded = self.save_and_load(model, f"initializer_{i}") + restored = loaded.layers[-1].kernel_initializer + self.assertIs(type(restored), type(init)) + np.testing.assert_array_equal(restored((4, 6)), init((4, 6))) From 96308345964322af234db722b419552c5cab8f83 Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Mon, 12 Oct 2026 16:52:29 +0200 Subject: [PATCH 10/16] docs: describe saving and loading models with complex layers Explain that models save to .keras and load without custom_objects, and that complexnn has to be imported before loading, because importing it is what registers the classes. Checked in a fresh process on Keras 3, Keras 2 and tf-keras: loading works after the import and fails without it. --- README.md | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/README.md b/README.md index eb32873..40f176b 100644 --- a/README.md +++ b/README.md @@ -53,6 +53,19 @@ model.compile(optimizer=optimizers.Adam(), loss='mse') An example working implementation of an autoencoder can be found [here](https://github.com/JesperDramsch/Complex-CNN-Seismic/). +Saving and loading +------------------ + +Models with complexnn layers save to the `.keras` format and load again without `custom_objects`. The layers are registered with Keras when `complexnn` is imported, so import it before loading: + +``` +import complexnn +import tensorflow as tf + +model.save("model.keras") +model = tf.keras.models.load_model("model.keras") +``` + Complex Format of Tensors ------------------------- From 75fa0be14211c38f47d3635c8c6968ea74ee6012 Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Mon, 12 Oct 2026 17:40:30 +0200 Subject: [PATCH 11/16] fix(dense): load ComplexDense with a non-default kernel initialiser ComplexDense.get_config stores any kernel initialiser other than "complex" as a dict. __init__ then tested `kernel_initializer in {"complex"}`, which hashes that dict and raised "unhashable type: 'dict'". A ComplexDense built with a Keras initialiser (by name or as an object) or with a complexnn initialiser such as ComplexInit could therefore not be rebuilt with from_config or loaded from a .keras file, on Keras 3, Keras 2 and tf-keras. Compare with == "complex" instead, in __init__, build and get_config. The saving tests gain ComplexDense entries with a seeded GlorotUniform and with a ComplexInit instance, so both the from_config and the .keras round trips cover a non-default kernel initialiser. A new test in test_dense checks that from_config restores the initialiser type for "complex", "glorot_uniform" and a ComplexInit instance. All six new subtests fail without the fix. --- complexnn/dense.py | 8 +++++--- tests/test_dense.py | 22 ++++++++++++++++++++++ tests/test_saving.py | 17 +++++++++++++++++ 3 files changed, 44 insertions(+), 3 deletions(-) diff --git a/complexnn/dense.py b/complexnn/dense.py index 003ef00..3bc0cbf 100644 --- a/complexnn/dense.py +++ b/complexnn/dense.py @@ -80,7 +80,9 @@ def __init__( self.activation = activations.get(activation) self.use_bias = use_bias self.init_criterion = init_criterion - if kernel_initializer in {"complex"}: + # get_config stores any other initialiser as a dict, which a set + # membership test would try to hash, so compare with == instead. + if kernel_initializer == "complex": self.kernel_initializer = kernel_initializer else: self.kernel_initializer = initializers.get(kernel_initializer) @@ -129,7 +131,7 @@ def init_w_real(shape, dtype=None): def init_w_imag(shape, dtype=None): return rng.normal(size=kernel_shape, loc=0, scale=s) # .astype(dtype) - if self.kernel_initializer in {"complex"}: + if self.kernel_initializer == "complex": real_init = init_w_real imag_init = init_w_imag else: @@ -198,7 +200,7 @@ def compute_output_shape(self, input_shape): return tuple(output_shape) def get_config(self): - if self.kernel_initializer in {"complex"}: + if self.kernel_initializer == "complex": ki = self.kernel_initializer else: ki = initializers.serialize(self.kernel_initializer) diff --git a/tests/test_dense.py b/tests/test_dense.py index b65ccc1..98721a9 100644 --- a/tests/test_dense.py +++ b/tests/test_dense.py @@ -18,3 +18,25 @@ def test_outputs_dense(self): true = (None, 16 * 2) calc = layer.compute_output_shape(input_shape) self.assertEqual(true, calc) + + +class TestDenseConfig(unittest.TestCase): + """ComplexDense configs rebuild the same kernel initialiser""" + + def test_kernel_initializer_round_trip(self): + """A kernel initialiser stored as a dict is read back, not hashed""" + cases = [ + ("complex", str), + ("glorot_uniform", tf.keras.initializers.GlorotUniform), + ( + conn.ComplexInit(kernel_size=(4, 4), input_dim=4, weight_dim=2), + conn.ComplexInit, + ), + ] + for kernel_initializer, expected_type in cases: + with self.subTest(expected_type.__name__): + layer = conn.ComplexDense(4, kernel_initializer=kernel_initializer) + config = layer.get_config() + clone = conn.ComplexDense.from_config(config) + self.assertIsInstance(clone.kernel_initializer, expected_type) + self.assertEqual(clone.get_config(), config) diff --git a/tests/test_saving.py b/tests/test_saving.py index 7e9b84d..1df7ced 100644 --- a/tests/test_saving.py +++ b/tests/test_saving.py @@ -13,6 +13,23 @@ # layer, so tests never share an instance. LAYERS = { "ComplexDense": ((8,), lambda: conn.ComplexDense(4)), + # A seeded Keras initialiser: Keras 2 warns when an unseeded one is called + # twice, and ComplexDense calls it for the real and the imaginary kernel. + "ComplexDense GlorotUniform": ( + (8,), + lambda: conn.ComplexDense( + 4, kernel_initializer=tf.keras.initializers.GlorotUniform(seed=1) + ), + ), + "ComplexDense ComplexInit": ( + (8,), + lambda: conn.ComplexDense( + 4, + kernel_initializer=conn.ComplexInit( + kernel_size=(4, 4), input_dim=4, weight_dim=2 + ), + ), + ), "ComplexConv1D": ((10, 4), lambda: conn.ComplexConv1D(2, 3, padding="same")), "ComplexConv2D": ((6, 6, 4), lambda: conn.ComplexConv2D(2, 3, padding="same")), "ComplexConv2D normalize_weight": ( From d166195ed96c6416416ed5e06e78ced0d92ba021 Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Mon, 12 Oct 2026 17:42:29 +0200 Subject: [PATCH 12/16] test(saving): check layer norm round trips on inputs that use gamma ComplexLayerNorm and LayerNormalization were round-tripped on a 2D input of shape (8,). With axis=-1 both layers then reduce over no axes, centre the input to exactly zero and return beta, whatever gamma holds. The prediction checks could not tell whether gamma was restored, so the comment that matching outputs show restored weights did not hold for these two entries. Use (4, 4, 8) for ComplexLayerNorm and (4, 8) for LayerNormalization. On these shapes a clone that gets beta but keeps its initial gamma gives a different output (checked on Keras 3 and Keras 2), whereas on (8,) it passed. The .keras round trip now also compares every loaded weight with the saved one, so restoring weights does not rest on the predictions alone. The 2D behaviour itself predates this branch and is left for a separate fix in norm.py. --- tests/test_saving.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/test_saving.py b/tests/test_saving.py index 1df7ced..fbdad18 100644 --- a/tests/test_saving.py +++ b/tests/test_saving.py @@ -42,8 +42,10 @@ ), "ComplexConv3D": ((4, 4, 4, 4), lambda: conn.ComplexConv3D(2, 3, padding="same")), "ComplexBN": ((4, 4, 8), lambda: conn.ComplexBN()), - "ComplexLayerNorm": ((8,), lambda: conn.ComplexLayerNorm()), - "LayerNormalization": ((8,), lambda: conn.LayerNormalization()), + # On a 2D input both layer norms reduce over no axes and return beta, so + # they get inputs with a feature axis on which the output depends on gamma. + "ComplexLayerNorm": ((4, 4, 8), lambda: conn.ComplexLayerNorm()), + "LayerNormalization": ((4, 8), lambda: conn.LayerNormalization()), "GetReal": ((8,), lambda: conn.GetReal()), "GetImag": ((4, 4, 8), lambda: conn.GetImag()), "GetAbs": ((8,), lambda: conn.GetAbs()), @@ -128,6 +130,10 @@ def test_layers(self): model = build_model(shape, make_layer, rng) loaded = self.save_and_load(model, f"layer_{i}") self.assertIs(type(loaded.layers[-1]), type(model.layers[-1])) + restored = loaded.get_weights() + self.assertEqual(len(restored), len(model.get_weights())) + for w_loaded, w_saved in zip(restored, model.get_weights()): + np.testing.assert_array_equal(w_loaded, w_saved) x = rng.standard_normal((3, *shape)).astype("float32") np.testing.assert_allclose( From ab5f82e98410135b48a566849ffdd742a2b68f65 Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Mon, 12 Oct 2026 17:53:37 +0200 Subject: [PATCH 13/16] test(saving): round-trip layers with non-default arguments TestLayerConfig and the .keras round trip built every layer with default arguments. A key that get_config left out was then missing on both sides of the config comparison, and from_config filled in the same default, so the check still passed. Removing epsilon (and momentum) from the ComplexBatchNormalization and ComplexLayerNorm configs went unnoticed by the whole suite, and the a84884a commit message promised more than the test checked. Build the layers with arguments away from their defaults: activation, strides, dilation_rate, use_bias, init_criterion, epsilon and momentum, spread over ComplexDense, the conv layers, ComplexBN, ComplexLayerNorm and LayerNormalization. A value that from_config does not restore now changes the config, and a lost argument that affects the output changes the output. Arguments that only affect initialisation or training (init_criterion, momentum) do not show in the output, so add a completeness test: every named __init__ argument of a working layer, and of the complexnn base class it passes arguments to, has to appear in get_config. The allow-list holds rank and ComplexConv1D's data_format, which the subclasses set themselves, and seed, which ComplexConv1D/2D/3D accept but do not use. The test compares for equality, so the allow-list has to shrink once a gap is closed. Removing epsilon from the BN, layer norm and conv configs, momentum from the BN config and activation from the dense config now fails 15 subtests on Keras 3, and each changed class fails at least one. --- tests/test_saving.py | 85 +++++++++++++++++++++++++++++++++++++++----- 1 file changed, 76 insertions(+), 9 deletions(-) diff --git a/tests/test_saving.py b/tests/test_saving.py index fbdad18..336517b 100644 --- a/tests/test_saving.py +++ b/tests/test_saving.py @@ -10,9 +10,16 @@ from tensorflow.keras.utils import get_registered_name, get_registered_object # Every working layer, with an input shape it accepts. Each entry builds a new -# layer, so tests never share an instance. +# layer, so tests never share an instance. Layers that take arguments set some +# of them away from their defaults: a value that from_config does not restore +# then changes the config, and a lost argument that affects the output +# (activation, strides, dilation_rate, epsilon) changes the output. Arguments +# missing from the config altogether are caught by TestConfigCompleteness. LAYERS = { - "ComplexDense": ((8,), lambda: conn.ComplexDense(4)), + "ComplexDense": ( + (8,), + lambda: conn.ComplexDense(4, activation="tanh", init_criterion="glorot"), + ), # A seeded Keras initialiser: Keras 2 warns when an unseeded one is called # twice, and ComplexDense calls it for the real and the imaginary kernel. "ComplexDense GlorotUniform": ( @@ -30,22 +37,39 @@ ), ), ), - "ComplexConv1D": ((10, 4), lambda: conn.ComplexConv1D(2, 3, padding="same")), - "ComplexConv2D": ((6, 6, 4), lambda: conn.ComplexConv2D(2, 3, padding="same")), + "ComplexConv1D": ( + (10, 4), + lambda: conn.ComplexConv1D( + 2, 3, padding="same", dilation_rate=2, activation="tanh" + ), + ), + "ComplexConv2D": ( + (6, 6, 4), + lambda: conn.ComplexConv2D( + 2, 3, strides=2, padding="same", init_criterion="glorot" + ), + ), "ComplexConv2D normalize_weight": ( (6, 6, 4), - lambda: conn.ComplexConv2D(2, 3, padding="same", normalize_weight=True), + lambda: conn.ComplexConv2D( + 2, 3, padding="same", normalize_weight=True, epsilon=1e-2 + ), ), "ComplexConv2D transposed": ( (6, 6, 4), lambda: conn.ComplexConv2D(1, 3, strides=2, padding="same", transposed=True), ), - "ComplexConv3D": ((4, 4, 4, 4), lambda: conn.ComplexConv3D(2, 3, padding="same")), - "ComplexBN": ((4, 4, 8), lambda: conn.ComplexBN()), + "ComplexConv3D": ( + (4, 4, 4, 4), + lambda: conn.ComplexConv3D( + 2, 3, padding="same", dilation_rate=2, use_bias=False + ), + ), + "ComplexBN": ((4, 4, 8), lambda: conn.ComplexBN(epsilon=1e-2, momentum=0.5)), # On a 2D input both layer norms reduce over no axes and return beta, so # they get inputs with a feature axis on which the output depends on gamma. - "ComplexLayerNorm": ((4, 4, 8), lambda: conn.ComplexLayerNorm()), - "LayerNormalization": ((4, 8), lambda: conn.LayerNormalization()), + "ComplexLayerNorm": ((4, 4, 8), lambda: conn.ComplexLayerNorm(epsilon=1e-2)), + "LayerNormalization": ((4, 8), lambda: conn.LayerNormalization(epsilon=1e-2)), "GetReal": ((8,), lambda: conn.GetReal()), "GetImag": ((4, 4, 8), lambda: conn.GetImag()), "GetAbs": ((8,), lambda: conn.GetAbs()), @@ -84,6 +108,49 @@ def test_from_config(self): ) +def init_parameters(cls): + """Named __init__ arguments of cls and of the complexnn classes it passes them to""" + names = set() + for klass in cls.__mro__: + if "__init__" in vars(klass) and klass.__module__.startswith("complexnn."): + for param in inspect.signature(klass.__init__).parameters.values(): + if param.name != "self" and param.kind not in ( + param.VAR_POSITIONAL, + param.VAR_KEYWORD, + ): + names.add(param.name) + return names + + +# Arguments that get_config leaves out. ComplexConv1D/2D/3D pass rank to +# ComplexConv themselves, and ComplexConv1D also data_format, so the config +# must not repeat them. They accept seed but never use it, so there is +# nothing to store yet. +NOT_IN_CONFIG = { + conn.ComplexConv1D: {"rank", "data_format", "seed"}, + conn.ComplexConv2D: {"rank", "seed"}, + conn.ComplexConv3D: {"rank", "seed"}, +} + + +class TestConfigCompleteness(unittest.TestCase): + """get_config stores every argument that __init__ takes""" + + def test_every_argument_in_config(self): + checked = set() + for _, make_layer in LAYERS.values(): + layer = make_layer() + cls = type(layer) + if cls in checked: + continue + checked.add(cls) + with self.subTest(cls.__name__): + missing = init_parameters(cls) - set(layer.get_config()) + # assertEqual rather than a subset check, so the allow-list has + # to shrink once a gap is closed. + self.assertEqual(missing, NOT_IN_CONFIG.get(cls, set())) + + class TestRegistration(unittest.TestCase): """Every class that complexnn exports is registered for serialisation""" From 48f522be05eabbff1205e35d8cab13e7b90f1910 Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Mon, 12 Oct 2026 17:56:42 +0200 Subject: [PATCH 14/16] fix(conv): keep the config of complexnn initialiser instances ComplexConv's sanitizedInitSer reduced any ComplexInit or ComplexIndependentFilters instance to the bare name "complex" or "complex_independent". from_config and clone_model then built a fresh initialiser from the layer's own init_criterion and the default seed, so a ComplexConv2D given ComplexInit(criterion="glorot", seed=5) came back with a "he" kernel drawn from seed 1337. Loading a .keras file still gave the same predictions, because the weights are restored, but the restored layer reported the wrong initialiser. Map only the names themselves to the names. Instances now go through initializers.serialize, which stores them as complexnn>ComplexInit or complexnn>ComplexIndependentFilters with their full config, and sanitizedInitGet already passes such dicts to initializers.get. Configs written before this change hold the names and still load. A new test in test_conv rebuilds a ComplexConv2D with each initialiser instance from its config and checks the initialiser type and config and the initial kernel. Both subtests fail without the fix. The saving tests gain a ComplexConv2D entry with a ComplexInit instance, and a .keras round trip restores ComplexInit with criterion "glorot" and seed 5 on Keras 3, Keras 2 and tf-keras. --- complexnn/conv.py | 9 +++++---- tests/test_conv.py | 39 +++++++++++++++++++++++++++++++++++++++ tests/test_saving.py | 11 +++++++++++ 3 files changed, 55 insertions(+), 4 deletions(-) diff --git a/complexnn/conv.py b/complexnn/conv.py index 923f265..5c7d8d3 100644 --- a/complexnn/conv.py +++ b/complexnn/conv.py @@ -172,10 +172,11 @@ def sanitizedInitSer(init): """sanitizedInitSer""" if init in [sqrt_init]: return "sqrt_init" - elif init == "complex" or isinstance(init, ComplexInit): - return "complex" - elif init == "complex_independent" or isinstance(init, ComplexIndependentFilters): - return "complex_independent" + elif init in ["complex", "complex_independent"]: + # Only these names stand for the initialisers that build() creates. An + # instance carries its own criterion and seed and is registered, so it + # goes through initializers.serialize like any other initialiser. + return init else: return initializers.serialize(init) diff --git a/tests/test_conv.py b/tests/test_conv.py index c83da31..108998d 100644 --- a/tests/test_conv.py +++ b/tests/test_conv.py @@ -76,3 +76,42 @@ def test_epsilon_in_config(self): ) clone = conn.ComplexConv2D.from_config(layer.get_config()) self.assertEqual(clone.epsilon, 1e-3) + + def test_complex_initializer_instance_in_config(self): + """A complexnn initialiser instance keeps its own config, seed included""" + initializers = [ + conn.ComplexInit( + kernel_size=(3, 3), + input_dim=2, + weight_dim=2, + nb_filters=2, + criterion="glorot", + seed=5, + ), + conn.ComplexIndependentFilters( + kernel_size=(3, 3), + input_dim=2, + weight_dim=2, + nb_filters=2, + criterion="he", + seed=5, + ), + ] + for init in initializers: + with self.subTest(type(init).__name__): + layer = conn.ComplexConv2D( + filters=2, kernel_size=3, kernel_initializer=init + ) + config = layer.get_config() + clone = conn.ComplexConv2D.from_config(config) + self.assertIs(type(clone.kernel_initializer), type(init)) + self.assertEqual( + clone.kernel_initializer.get_config(), init.get_config() + ) + self.assertEqual(clone.get_config(), config) + + layer.build((None, 6, 6, 4)) + clone.build((None, 6, 6, 4)) + np.testing.assert_array_equal( + np.asarray(clone.kernel), np.asarray(layer.kernel) + ) diff --git a/tests/test_saving.py b/tests/test_saving.py index 336517b..9aa84dc 100644 --- a/tests/test_saving.py +++ b/tests/test_saving.py @@ -55,6 +55,17 @@ 2, 3, padding="same", normalize_weight=True, epsilon=1e-2 ), ), + "ComplexConv2D ComplexInit": ( + (6, 6, 4), + lambda: conn.ComplexConv2D( + 2, + 3, + padding="same", + kernel_initializer=conn.ComplexInit( + kernel_size=(3, 3), input_dim=2, weight_dim=2, nb_filters=2, seed=5 + ), + ), + ), "ComplexConv2D transposed": ( (6, 6, 4), lambda: conn.ComplexConv2D(1, 3, strides=2, padding="same", transposed=True), From 22daa5da5b618e3739bfd0006c3a8eb7c88d5c24 Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Mon, 12 Oct 2026 19:49:01 +0200 Subject: [PATCH 15/16] fix(saving): load complexnn initialisers from legacy .h5 files The .h5 loader looks a serialised initialiser up by its bare class name, while complexnn registers its classes as complexnn>ClassName, so a ComplexDense or ComplexConv with a ComplexInit instance no longer loaded once the full config was stored. get_initializer falls back to the registered name when Keras cannot resolve such a config. A regression test round-trips both through .h5 on Keras 3, Keras 2 and tf-keras. --- complexnn/bn.py | 4 +++- complexnn/conv.py | 3 ++- complexnn/dense.py | 6 +++--- complexnn/utils.py | 23 ++++++++++++++++++++++- tests/test_saving.py | 33 +++++++++++++++++++++++++++++++++ 5 files changed, 63 insertions(+), 6 deletions(-) diff --git a/complexnn/bn.py b/complexnn/bn.py index 0ad5195..2f3f90c 100644 --- a/complexnn/bn.py +++ b/complexnn/bn.py @@ -12,6 +12,8 @@ from tensorflow.keras.layers import InputSpec, Layer from tensorflow.keras.utils import register_keras_serializable +from .utils import get_initializer + def sqrt_init(shape, dtype=None): value = (1 / np.sqrt(2)) * K.ones(shape) @@ -22,7 +24,7 @@ def sanitizedInitGet(init): if init in ["sqrt_init"]: return sqrt_init else: - return initializers.get(init) + return get_initializer(init) def sanitizedInitSer(init): diff --git a/complexnn/conv.py b/complexnn/conv.py index 5c7d8d3..5d61017 100644 --- a/complexnn/conv.py +++ b/complexnn/conv.py @@ -17,6 +17,7 @@ from .bn import sqrt_init from .fft import fft, fft2, ifft, ifft2 from .init import ComplexIndependentFilters, ComplexInit +from .utils import get_initializer def conv1d_transpose( @@ -165,7 +166,7 @@ def sanitizedInitGet(init): elif init in ["complex", "complex_independent", "glorot_complex", "he_complex"]: return init else: - return initializers.get(init) + return get_initializer(init) def sanitizedInitSer(init): diff --git a/complexnn/dense.py b/complexnn/dense.py index 3bc0cbf..bd03dd7 100644 --- a/complexnn/dense.py +++ b/complexnn/dense.py @@ -8,7 +8,7 @@ from tensorflow.keras.layers import InputSpec, Layer from tensorflow.keras.utils import register_keras_serializable -from .utils import _compute_fans +from .utils import _compute_fans, get_initializer @register_keras_serializable(package="complexnn") @@ -85,8 +85,8 @@ def __init__( if kernel_initializer == "complex": self.kernel_initializer = kernel_initializer else: - self.kernel_initializer = initializers.get(kernel_initializer) - self.bias_initializer = initializers.get(bias_initializer) + self.kernel_initializer = get_initializer(kernel_initializer) + self.bias_initializer = get_initializer(bias_initializer) self.kernel_regularizer = regularizers.get(kernel_regularizer) self.bias_regularizer = regularizers.get(bias_regularizer) self.activity_regularizer = regularizers.get(activity_regularizer) diff --git a/complexnn/utils.py b/complexnn/utils.py index a7ffad3..d86dc44 100644 --- a/complexnn/utils.py +++ b/complexnn/utils.py @@ -3,8 +3,29 @@ import numpy as np import tensorflow.keras.backend as K +from tensorflow.keras import initializers from tensorflow.keras.layers import Lambda, Layer -from tensorflow.keras.utils import register_keras_serializable +from tensorflow.keras.utils import get_registered_object, register_keras_serializable + + +def get_initializer(identifier): + """Resolve an initialiser like ``initializers.get``, complexnn ones included. + + The legacy .h5 loader looks a serialised initialiser up by its bare class + name, while complexnn registers its classes as ``complexnn>ClassName``, so + a ``ComplexInit`` saved with its config would not load. When Keras cannot + resolve such a config, look it up by its registered name instead. + """ + try: + return initializers.get(identifier) + except (ValueError, TypeError): + if isinstance(identifier, dict): + registered = str(identifier.get("registered_name") or "") + cls = get_registered_object(registered) if registered else None + if cls is not None and registered.startswith("complexnn>"): + return cls.from_config(identifier.get("config") or {}) + raise + # # GetReal/GetImag Lambda layer Implementation diff --git a/tests/test_saving.py b/tests/test_saving.py index 9aa84dc..58a191e 100644 --- a/tests/test_saving.py +++ b/tests/test_saving.py @@ -1,4 +1,5 @@ import inspect +import json import os import tempfile import unittest @@ -221,6 +222,38 @@ def test_layers(self): atol=1e-6, ) + def test_h5_with_complex_initializer(self): + """Legacy .h5 files keep a complexnn initialiser instance + + Keras 2's .h5 loader looks an initialiser up by its bare class name, + which is not where complexnn registers it. + """ + rng = np.random.default_rng(2) + for name in ("ComplexDense ComplexInit", "ComplexConv2D ComplexInit"): + shape, make_layer = LAYERS[name] + with self.subTest(name): + model = build_model(shape, make_layer, rng) + path = os.path.join(self.tmp, name.replace(" ", "_") + ".h5") + with warnings.catch_warnings(): + # Keras 3 warns that .h5 is a legacy format. + warnings.simplefilter("ignore") + model.save(path) + loaded = tf.keras.models.load_model(path) + restored = loaded.layers[-1].kernel_initializer + self.assertIs(type(restored), conn.ComplexInit) + # The file stores JSON, which turns tuples into lists. + saved_config = model.layers[-1].kernel_initializer.get_config() + self.assertEqual( + restored.get_config(), json.loads(json.dumps(saved_config)) + ) + x = rng.standard_normal((3, *shape)).astype("float32") + np.testing.assert_allclose( + loaded.predict(x, verbose=0), + model.predict(x, verbose=0), + rtol=1e-6, + atol=1e-6, + ) + def test_initializers(self): """A built-in layer keeps a complex initialiser through saving and loading""" initializers = [ From 4eba825fd6d426457d5aa8fc8097c31768f934b7 Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Mon, 12 Oct 2026 20:19:34 +0200 Subject: [PATCH 16/16] test(saving): compare restored weights with zip(strict=True) --- tests/test_saving.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_saving.py b/tests/test_saving.py index 58a191e..627c6f3 100644 --- a/tests/test_saving.py +++ b/tests/test_saving.py @@ -211,7 +211,9 @@ def test_layers(self): self.assertIs(type(loaded.layers[-1]), type(model.layers[-1])) restored = loaded.get_weights() self.assertEqual(len(restored), len(model.get_weights())) - for w_loaded, w_saved in zip(restored, model.get_weights()): + for w_loaded, w_saved in zip( + restored, model.get_weights(), strict=True + ): np.testing.assert_array_equal(w_loaded, w_saved) x = rng.standard_normal((3, *shape)).astype("float32")