Fix Keras build/call parameter naming in DeepSet layers

[SermetPekin]

Fix for issue

The error originates from this line in the Keras source code

# old signratures 
 @sanitize_input_shape
    def build(self, input_shape):
        self.call(keras.ops.zeros(input_shape))

    def call(self, input_set: Tensor, training: bool = False, **kwargs) -> Tensor:
        """Performs the forward pass of a learnable equivariant transform.

# new signatures
@sanitize_input_shape
    def build(self, input_set_shape):
        self.call(keras.ops.zeros(input_set_shape))

    def call(self, input_set: Tensor, training: bool = False, **kwargs) -> Tensor:
        """Performs the forward pass of a learnable equivariant transform.

On Windows with python 3.10.8
I get the error below.

FAILED tests/test_approximators/test_approximator_standardization/
test_approximator_standardization.py::test_save_and_load[standardize='all'] 
- ValueError: For a `build()` method with more than one argument, all arguments should have a `_shape` suffix
 and match an argument from `call()`. E.g. `build(self, foo_shape, bar_shape)`  
For layer 'EquivariantLayer', Received `build()` argument `self`, which does not end in `_shape`.

As I previously mentioned I got this error on Windows (which I could not reproduce with same packages on mac or ubuntu)
I created environment using

uv venv --python 3.10
.venv\Scripts\activate
uv pip install ".[all]"
uv run pytest -m "not slow"  

I checked with Mac and Ubuntu they pass successfully. (uv run pytest -m 'not slow')

absl-py==2.3.1
accessible-pygments==0.0.5
alabaster==1.0.0
anyio==4.11.0
argon2-cffi==25.1.0
argon2-cffi-bindings==25.1.0
arrow==1.4.0
asttokens==3.0.0
astunparse==1.6.3
async-lru==2.0.5
attrs==25.4.0
babel==2.17.0
bayesflow @ file:///C:/Users/Username/PycharmProjects/bf_orig/bayesflow-main
beautifulsoup4==4.14.2
bleach==6.2.0
cachetools==6.2.1
certifi==2025.10.5
cffi==2.0.0
cfgv==3.4.0
chardet==5.2.0
charset-normalizer==3.4.4
click==8.3.0
colorama==0.4.6
comm==0.2.3
contourpy==1.3.2
coverage==7.11.0
cycler==0.12.1
debugpy==1.8.17
decorator==5.2.1
defusedxml==0.7.1
distlib==0.4.0
docutils==0.21.2
exceptiongroup==1.3.0
execnet==2.1.1
executing==2.2.1
fastjsonschema==2.21.2
filelock==3.20.0
flatbuffers==25.9.23
fonttools==4.60.1
fqdn==1.5.1
gast==0.6.0
google-pasta==0.2.0
greenlet==3.2.4
grpcio==1.76.0
h11==0.16.0
h5py==3.15.1
httpcore==1.0.9
httpx==0.28.1
identify==2.6.15
idna==3.11
imagesize==1.4.1
importlib-metadata==8.7.0
iniconfig==2.3.0
ipykernel==7.0.1
ipython==8.37.0
ipywidgets==8.1.7
isoduration==20.11.0
jedi==0.19.2
jinja2==3.1.6
json5==0.12.1
jsonpointer==3.0.0
jsonschema==4.25.1
jsonschema-specifications==2025.9.1
jupyter==1.1.1
jupyter-cache==1.0.1
jupyter-client==8.6.3
jupyter-console==6.6.3
jupyter-core==5.9.1
jupyter-events==0.12.0
jupyter-lsp==2.3.0
jupyter-server==2.17.0
jupyter-server-terminals==0.5.3
jupyterlab==4.4.10
jupyterlab-pygments==0.3.0
jupyterlab-server==2.28.0
jupyterlab-widgets==3.0.15
keras==3.11.3
kiwisolver==1.4.9
lark==1.3.0
latexcodec==3.0.1
libclang==18.1.1
line-profiler==5.0.0
markdown==3.9
markdown-it-py==3.0.0
markupsafe==3.0.3
matplotlib==3.10.7
matplotlib-inline==0.2.1
mdit-py-plugins==0.5.0
mdurl==0.1.2
mistune==3.1.4
ml-dtypes==0.5.3
myst-nb==1.3.0
myst-parser==4.0.1
namex==0.1.0
nbclient==0.10.2
nbconvert==7.16.6
nbformat==5.10.4
nest-asyncio==1.6.0
nodeenv==1.9.1
notebook==7.4.7
notebook-shim==0.2.4
numpy==1.26.4
numpydoc==1.9.0
opt-einsum==3.4.0
optree==0.17.0
overrides==7.7.0
packaging==25.0
pandas==2.3.3
pandocfilters==1.5.1
parso==0.8.5
pillow==12.0.0
pip==22.2.2
platformdirs==4.5.0
pluggy==1.6.0
pre-commit==4.3.0
prometheus-client==0.23.1
prompt-toolkit==3.0.52
protobuf==6.33.0
psutil==7.1.1
pure-eval==0.2.3
pybtex==0.25.1
pybtex-docutils==1.0.3
pycparser==2.23
pydata-sphinx-theme==0.16.1
pygments==2.19.2
pyparsing==3.2.5
pyproject-api==1.10.0
pytest==8.4.2
pytest-cov==7.0.0
pytest-rerunfailures==16.1
pytest-xdist==3.8.0
python-dateutil==2.9.0.post0
python-json-logger==4.0.0
pytz==2025.2
pywinpty==3.0.2
pyyaml==6.0.3
pyzmq==27.1.0
referencing==0.37.0
requests==2.32.5
rfc3339-validator==0.1.4
rfc3986-validator==0.1.1
rfc3987-syntax==1.1.0
rich==14.2.0
rpds-py==0.28.0
ruff==0.14.2
scipy==1.15.3
seaborn==0.13.2
send2trash==1.8.3
setuptools==63.2.0
six==1.17.0
sniffio==1.3.1
snowballstemmer==2.2.0
soupsieve==2.8
sphinx==8.1.3
sphinx-design==0.6.1
sphinx-polyversion==1.1.0
sphinxcontrib-applehelp==2.0.0
sphinxcontrib-bibtex==2.6.5
sphinxcontrib-devhelp==2.0.0
sphinxcontrib-htmlhelp==2.1.0
sphinxcontrib-jsmath==1.0.1
sphinxcontrib-qthelp==2.0.0
sphinxcontrib-serializinghtml==2.0.0
sqlalchemy==2.0.44
stack-data==0.6.3
tabulate==0.9.0
tensorboard==2.20.0
tensorboard-data-server==0.7.2
tensorflow==2.20.0
termcolor==3.1.0
terminado==0.18.1
tinycss2==1.4.0
tomli==2.3.0
tornado==6.5.2
tox==4.31.0
tqdm==4.67.1
traitlets==5.14.3
typing-extensions==4.15.0
tzdata==2025.2
uri-template==1.3.0
urllib3==2.5.0
irtualenv==20.35.3
wcwidth==0.2.14
webcolors==24.11.1
webencodings==0.5.1
websocket-client==1.9.0
werkzeug==3.1.3
wheel==0.45.1
widgetsnbextension==4.0.14
wrapt==2.0.0
zipp==3.23.0

[stefanradev93]

Hi Sermet, thanks for the first contribution! Can you provide some insight into the failure, as I am unable to reproduce it on my windows machine and do not have any intuition about why the fix works. The error msg is:

- ValueError: For a build()method with more than one argument, all arguments should have a_shape suffix, but both input_shape and input_set_shape have the required suffix.

[LarsKue]

Could be a new upstream issue. @stefanradev93 which version of keras did you test this on? The windows workflow did not catch this issue, but it has been a while since it ran on main.

The current code does not seem faulty, since the build method only has one argument, unlike what the error message describes. So both input_shape and input_set_shape argument names should be valid, regardless of their suffix.

[stefanradev93]

I am on 3.10.0.

[LarsKue]

@SermetPekin is on 3.11.3, so this could be it. I will investigate and report back shortly.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
bayesflow/networks/deep_set/equivariant_layer.py	100.00% <100.00%> (ø)
bayesflow/networks/deep_set/invariant_layer.py	100.00% <100.00%> (ø)

[SermetPekin]

Hi Sermet, thanks for the first contribution! Can you provide some insight into the failure, as I am unable to reproduce it on my windows machine and do not have any intuition about why the fix works. The error msg is:

- ValueError: For a build()method with more than one argument, all arguments should have a_shape suffix, but both input_shape and input_set_shape have the required suffix.

Thanks for the review @stefanradev93 .
I think the reason looks like keras at some point decided to force specific name format of these classes [ for methods call and build ] if it is x for call it should be x_shape for the other. it is just internal name change which would not make anything worse for backward compad (as long as those functions are passing positional ) I will test with 3.10.0 now.

[LarsKue]

@SermetPekin your assumption is close, but not quite true (unless they changed it recently). If my memory serves right, Keras assumes the following signature for the build and call methods:

def build(self, build_arg):
    ...

def call(self, call_arg):
    ...

where build_arg and call_arg can both be named anything. build_arg is then passed the value of keras.ops.shape(call_arg). However, this is incompatible with methods that have multiple call_args:

def call(self, x, y):
    # e.g. a Transformer
    ...

Their solution is to then require that the build method specifies in its signature which shape arguments it needs. For the example above, this then becomes

def build(self, x_shape, y_shape):
    ...

where the naming of the arguments is now fixed.

In our case, this should not be a problem, since we only have one forward argument Tensor.

[SermetPekin]

is it possible that it is due to we have 2 _ in input_set_shape ? (just brain storming )
When I search in keras code this is throwing it
keras

# keras 3.10.0 + tensorflow 2.20.0 

pytest -v tests/test_approximators/test_approximator_standardization/test_approximator_standardization.py::test_save_and_load 

FAILED tests/test_approximators/test_approximator_standardization/test_approximator_standardization.py::test_save_and_load[standardize='all'] - ValueError: For a `build()` method with more than one argument, all arguments should have a `_shape` suffix and match an argument from `call()`. E.g. `build(self, foo_shape, bar_shape)`  For layer 'EquivariantLayer', Received `build()` argument `self`, which does not end in `_shape`.
FAILED tests/test_approximators/test_approximator_standardization/test_approximator_standardization.py::test_save_and_load[standardize=None] - ValueError: For a `build()` method with more than one argument, all arguments should have a `_shape` suffix and match an argument from `call()`. E.g. `build(self, foo_shape, bar_shape)`  For layer 'EquivariantLayer', Received `build()` argument `self`, which does not end in `_shape`.
FAILED tests/test_approximators/test_approximator_standardization/test_approximator_standardization.py::test_save_and_load[standardize='inference_variables'] - ValueError: For a `build()` method with more than one argument, all arguments should have a `_shape` suffix and match an argument from `call()`. E.g. `build(self, foo_shape, bar_shape)`  For layer 'EquivariantLayer', Received `build()` argument `self`, which does not end in `_shape`.
FAILED tests/test_approximators/test_approximator_standardization/test_approximator_standardization.py::test_save_and_load[standardize='summary_variables'] - ValueError: For a `build()` method with more than one argument, all arguments should have a `_shape` suffix and match an argument from `call()`. E.g. `build(self, foo_shape, bar_shape)`  For layer 'EquivariantLayer', Received `build()` argument `self`, which does not end in `_shape`.
FAILED tests/test_approximators/test_approximator_standardization/test_approximator_standardization.py::test_save_and_load[standardize=('inference_variables', 'summary_variables', 'inference_conditions')] - ValueError: For a `build()` method with more than one argument, all arguments should have a `_shape` suffix and match an argument from `call()`. E.g. `build(self, foo_shape, bar_shape)`  For layer 'EquivariantLayer', Received `build()` argument `self`, which does not end in `_shape`.

[LarsKue]

In other news, I re-ran the workflow and could not reproduce the issue, even with version 3.11.3, so marking this as not reproducible for now. Below, you may find the package versions that I tested this with, so if you keep encountering the issue and can reliably reproduce it, please re-open this PR:

Windows + JAX

Package                   Version
------------------------- -----------
absl-py                   2.3.1
asttokens                 3.0.0
attrs                     25.4.0
bayesflow                 2.0.7
beautifulsoup4            4.14.2
bleach                    6.2.0
colorama                  0.4.6
comm                      0.2.3
contourpy                 1.3.2
coverage                  7.11.0
cycler                    0.12.1
debugpy                   1.8.17
decorator                 5.2.1
defusedxml                0.7.1
exceptiongroup            1.3.0
execnet                   2.1.1
executing                 2.2.1
fastjsonschema            2.21.2
fonttools                 4.60.1
h5py                      3.15.1
iniconfig                 2.3.0
ipykernel                 7.1.0
ipython                   8.37.0
jax                       0.6.2
jaxlib                    0.6.2
jedi                      0.19.2
Jinja2                    3.1.6
jsonschema                4.25.1
jsonschema-specifications 2025.9.1
jupyter_client            8.6.3
jupyter_core              5.9.1
jupyterlab_pygments       0.3.0
keras                     3.11.3
kiwisolver                1.4.9
markdown-it-py            4.0.0
MarkupSafe                3.0.3
matplotlib                3.10.7
matplotlib-inline         0.2.1
mdurl                     0.1.2
mistune                   3.1.4
ml_dtypes                 0.5.3
namex                     0.1.0
nbclient                  0.10.2
nbconvert                 7.16.6
nbformat                  5.10.4
nest-asyncio              1.6.0
numpy                     1.26.4
opt_einsum                3.4.0
optree                    0.17.0
packaging                 25.0
pandas                    2.3.3
pandocfilters             1.5.1
parso                     0.8.5
pillow                    12.0.0
pip                       25.3
platformdirs              4.5.0
pluggy                    1.6.0
prompt_toolkit            3.0.52
psutil                    7.1.2
pure_eval                 0.2.3
Pygments                  2.19.2
pyparsing                 3.2.5
pytest                    8.4.2
pytest-cov                7.0.0
pytest-rerunfailures      16.1
pytest-xdist              3.8.0
python-dateutil           2.9.0.post0
pytz                      2025.2
pyzmq                     27.1.0
referencing               0.37.0
rich                      14.2.0
rpds-py                   0.28.0
scipy                     1.15.3
seaborn                   0.13.2
setuptools                80.9.0
six                       1.17.0
soupsieve                 2.8
stack-data                0.6.3
tinycss2                  1.4.0
tomli                     2.3.0
tornado                   6.5.2
tqdm                      4.67.1
traitlets                 5.14.3
typing_extensions         4.15.0
tzdata                    2025.2
wcwidth                   0.2.14
webencodings              0.5.1
wheel                     0.45.1

Windows + TensorFlow

Package                   Version
------------------------- -----------
absl-py                   2.3.1
asttokens                 3.0.0
astunparse                1.6.3
attrs                     25.4.0
bayesflow                 2.0.7
beautifulsoup4            4.14.2
bleach                    6.2.0
certifi                   2025.10.5
charset-normalizer        3.4.4
colorama                  0.4.6
comm                      0.2.3
contourpy                 1.3.2
coverage                  7.11.0
cycler                    0.12.1
debugpy                   1.8.17
decorator                 5.2.1
defusedxml                0.7.1
exceptiongroup            1.3.0
execnet                   2.1.1
executing                 2.2.1
fastjsonschema            2.21.2
flatbuffers               25.9.23
fonttools                 4.60.1
gast                      0.6.0
google-pasta              0.2.0
grpcio                    1.76.0
h5py                      3.15.1
idna                      3.11
iniconfig                 2.3.0
ipykernel                 7.1.0
ipython                   8.37.0
jedi                      0.19.2
Jinja2                    3.1.6
jsonschema                4.25.1
jsonschema-specifications 2025.9.1
jupyter_client            8.6.3
jupyter_core              5.9.1
jupyterlab_pygments       0.3.0
keras                     3.11.3
kiwisolver                1.4.9
libclang                  18.1.1
Markdown                  3.9
markdown-it-py            4.0.0
MarkupSafe                3.0.3
matplotlib                3.10.7
matplotlib-inline         0.2.1
mdurl                     0.1.2
mistune                   3.1.4
ml_dtypes                 0.5.3
namex                     0.1.0
nbclient                  0.10.2
nbconvert                 7.16.6
nbformat                  5.10.4
nest-asyncio              1.6.0
numpy                     1.26.4
opt_einsum                3.4.0
optree                    0.17.0
packaging                 25.0
pandas                    2.3.3
pandocfilters             1.5.1
parso                     0.8.5
pillow                    12.0.0
pip                       25.3
platformdirs              4.5.0
pluggy                    1.6.0
prompt_toolkit            3.0.52
protobuf                  6.33.0
psutil                    7.1.2
pure_eval                 0.2.3
Pygments                  2.19.2
pyparsing                 3.2.5
pytest                    8.4.2
pytest-cov                7.0.0
pytest-rerunfailures      16.1
pytest-xdist              3.8.0
python-dateutil           2.9.0.post0
pytz                      2025.2
pyzmq                     27.1.0
referencing               0.37.0
requests                  2.32.5
rich                      14.2.0
rpds-py                   0.28.0
scipy                     1.15.3
seaborn                   0.13.2
setuptools                80.9.0
six                       1.17.0
soupsieve                 2.8
stack-data                0.6.3
tensorboard               2.20.0
tensorboard-data-server   0.7.2
tensorflow                2.20.0
termcolor                 3.2.0
tinycss2                  1.4.0
tomli                     2.3.0
tornado                   6.5.2
tqdm                      4.67.1
traitlets                 5.14.3
typing_extensions         4.15.0
tzdata                    2025.2
urllib3                   2.5.0
wcwidth                   0.2.14
webencodings              0.5.1
Werkzeug                  3.1.3
wheel                     0.45.1
wrapt                     2.0.0

Windows + PyTorch

Package                   Version
------------------------- -----------
absl-py                   2.3.1
asttokens                 3.0.0
attrs                     25.4.0
bayesflow                 2.0.7
beautifulsoup4            4.14.2
bleach                    6.2.0
colorama                  0.4.6
comm                      0.2.3
contourpy                 1.3.2
coverage                  7.11.0
cycler                    0.12.1
debugpy                   1.8.17
decorator                 5.2.1
defusedxml                0.7.1
exceptiongroup            1.3.0
execnet                   2.1.1
executing                 2.2.1
fastjsonschema            2.21.2
filelock                  3.19.1
fonttools                 4.60.1
fsspec                    2025.9.0
h5py                      3.15.1
iniconfig                 2.3.0
ipykernel                 7.1.0
ipython                   8.37.0
jedi                      0.19.2
Jinja2                    3.1.6
jsonschema                4.25.1
jsonschema-specifications 2025.9.1
jupyter_client            8.6.3
jupyter_core              5.9.1
jupyterlab_pygments       0.3.0
keras                     3.11.3
kiwisolver                1.4.9
markdown-it-py            4.0.0
MarkupSafe                3.0.3
matplotlib                3.10.7
matplotlib-inline         0.2.1
mdurl                     0.1.2
mistune                   3.1.4
ml_dtypes                 0.5.3
mpmath                    1.3.0
namex                     0.1.0
nbclient                  0.10.2
nbconvert                 7.16.6
nbformat                  5.10.4
nest-asyncio              1.6.0
networkx                  3.3
numpy                     1.26.4
optree                    0.17.0
packaging                 25.0
pandas                    2.3.3
pandocfilters             1.5.1
parso                     0.8.5
pillow                    12.0.0
pip                       25.3
platformdirs              4.5.0
pluggy                    1.6.0
prompt_toolkit            3.0.52
psutil                    7.1.2
pure_eval                 0.2.3
Pygments                  2.19.2
pyparsing                 3.2.5
pytest                    8.4.2
pytest-cov                7.0.0
pytest-rerunfailures      16.1
pytest-xdist              3.8.0
python-dateutil           2.9.0.post0
pytz                      2025.2
pyzmq                     27.1.0
referencing               0.37.0
rich                      14.2.0
rpds-py                   0.28.0
scipy                     1.15.3
seaborn                   0.13.2
setuptools                80.9.0
six                       1.17.0
soupsieve                 2.8
stack-data                0.6.3
sympy                     1.14.0
tinycss2                  1.4.0
tomli                     2.3.0
torch                     2.9.0+cpu
torchaudio                2.9.0+cpu
torchvision               0.24.0+cpu
tornado                   6.5.2
tqdm                      4.67.1
traitlets                 5.14.3
typing_extensions         4.15.0
tzdata                    2025.2
wcwidth                   0.2.14
webencodings              0.5.1
wheel                     0.45.1

[LarsKue]

@SermetPekin I don't think it is likely, but feel free to try using inputset and inputset_shape instead.

[SermetPekin]

can it be due to sanitize_input_shape decorator on top of build ?

[SermetPekin]

@LarsKue , @stefanradev93 I will check this issue and come back with a detailed report soon. Thanks for your valuable feedback and your time.