[Failing Test]: Beam ML unit tests fail with Keras 3.3.3 and TF 2.16.1

[tvalentyn]

What happened?

Sample failure:

$ pytest apache_beam/ml/inference/tensorflow_inference_test.py::TFRunInferenceTest::test_predict_numpy_with_batch_size
================================================================================================================================== test session starts ===================================================================================================================================
platform linux -- Python 3.9.18, pytest-7.4.4, pluggy-1.5.0
rootdir: /home/valentyn/projects/beam/beam/beam/sdks/python
configfile: pytest.ini
plugins: xdist-3.6.1, requests-mock-1.12.1, timeout-2.3.1, hypothesis-6.102.1
timeout: 600.0s
timeout method: signal
timeout func_only: False
collected 1 item                                                                                                                                                                                                                                                                         

apache_beam/ml/inference/tensorflow_inference_test.py F                                                                                                                                                                                                                            [100%]

======================================================================================================================================== FAILURES ========================================================================================================================================
_________________________________________________________________________________________________________________ TFRunInferenceTest.test_predict_numpy_with_batch_size __________________________________________________________________________________________________________________

self = <apache_beam.ml.inference.tensorflow_inference_test.TFRunInferenceTest testMethod=test_predict_numpy_with_batch_size>

    def test_predict_numpy_with_batch_size(self):
>     model = _create_mult2_model()

apache_beam/ml/inference/tensorflow_inference_test.py:220: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
apache_beam/ml/inference/tensorflow_inference_test.py:68: in _create_mult2_model
    inputs = tf.keras.Input(shape=(3))
test-suites/tox/py39/build/srcs/sdks/python/target/.tox-py39-ml/py39-ml/lib/python3.9/site-packages/keras/src/layers/core/input_layer.py:143: in Input
    layer = InputLayer(
test-suites/tox/py39/build/srcs/sdks/python/target/.tox-py39-ml/py39-ml/lib/python3.9/site-packages/keras/src/layers/core/input_layer.py:46: in __init__
    shape = backend.standardize_shape(shape)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

shape = 3

    def standardize_shape(shape):
        if not isinstance(shape, tuple):
            if shape is None:
                raise ValueError("Undefined shapes are not supported.")
            if not hasattr(shape, "__iter__"):
>               raise ValueError(f"Cannot convert '{shape}' to a shape.")
E               ValueError: Cannot convert '3' to a shape.

test-suites/tox/py39/build/srcs/sdks/python/target/.tox-py39-ml/py39-ml/lib/python3.9/site-packages/keras/src/backend/common/variables.py:530: ValueError

Issue Failure

Failure: Test is continually failing

Issue Priority

Priority: 2 (backlog / disabled test but we think the product is healthy)

Issue Components

• [X] Component: Python SDK
• [ ] Component: Java SDK
• [ ] Component: Go SDK
• [ ] Component: Typescript SDK
• [ ] Component: IO connector
• [ ] Component: Beam YAML
• [ ] Component: Beam examples
• [ ] Component: Beam playground
• [ ] Component: Beam katas
• [ ] Component: Website
• [ ] Component: Spark Runner
• [ ] Component: Flink Runner
• [ ] Component: Samza Runner
• [ ] Component: Twister2 Runner
• [ ] Component: Hazelcast Jet Runner
• [ ] Component: Google Cloud Dataflow Runner

[tvalentyn]

cc: @damccorm - we might want to prioritize https://github.com/apache/beam/issues/30908.

[tvalentyn]

to repro:

gradlew :sdks:python:test-suites:tox:py39:testPy39ML -Pposargs=apache_beam/ml/