Repository navigation
feat: save and load models with complex layers - #31
Merged
Merged
Conversation
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.
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.
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.
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.
…on 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.
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.
…hape 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.
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.
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.
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.
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.
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.
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.
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.
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.
# Conflicts: # complexnn/conv.py # complexnn/dense.py # complexnn/init.py # complexnn/norm.py # complexnn/pool.py # complexnn/utils.py
This was referenced Oct 12, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stage 2, PR 2 of 4. Models with complex layers save to
.kerasand.h5and load withoutcustom_objects.What changes
register_keras_serializable(package="complexnn"), so loading needs nocustom_objects. A test checks that every exported class is registered.ComplexConv1Dstoreddata_format, which its constructor rejects.ComplexConvlostepsilon.ComplexLayerNormstored itssqrt_initdefault as a bare function.LayerNormalizationdropped its initialisers and regularisers.ComplexDensecrashed on a non-defaultkernel_initializer, whichfrom_confighands back as a dict.ComplexConvwas reduced to the name"complex". It now keeps its full config, seed included. Old configs that hold the names still load.get_configforComplexInit, andcompute_output_shapeforComplexBatchNormalizationandComplexLayerNorm..h5. The.h5loader looks an initialiser up by its bare class name, while complexnn registers it ascomplexnn>ClassName. Without help, a layer with aComplexInitinstance would not load.utils.get_initializerfalls back to the registered name.Tests
get_configstores every constructor argument, with a small allow-list forrank,data_formatand an unusedseed..kerasround trip: for every setup, with exact weight comparison..h5round trip: the twoComplexInitsetups. This fails on all three environments without the fix.Verification
input_shapeones, removed by the public-API PR)ruff format --check, pre-commit and the flake8 critical check are clean.Compatibility
.kerasfiles written by 0.3.0 load with this PR exactly as they load with 0.3.0. I checked this on Keras 3 with files saved by main:ComplexDense,ComplexConv2DorComplexBNneedcustom_objects.ComplexConv1DorComplexLayerNormfail even withcustom_objects. They fail on 0.3.0 itself too, because of the storeddata_formatandsqrt_initthis PR fixes.Files saved after this PR load with no
custom_objects.Merge notes
Merge after the public-API PR. The rebase is import lines only: both PRs edit the same import blocks.