Regression: Cannot export model with bilinear UpSampling2D and dynamic inputs

[pwuertz]

Since Keras3 (tested with v3.3.3), UpSampling2D(interpolation="bilinear") with dynamic input sizes causes model exports to fail. For example,

import keras as k

x_in = k.layers.Input(shape=[None, None, 1], name="input")
x_out = k.layers.UpSampling2D(interpolation="bilinear")(x_in)
model = k.Model(inputs=x_in, outputs=x_out)
model.export("./test_model_export")

raises the following exception on .export():

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/pwuertz/.local/lib/python3.11/site-packages/keras/src/models/model.py", line 514, in export
    export_lib.export_model(self, filepath)
  File "/home/pwuertz/.local/lib/python3.11/site-packages/keras/src/export/export_lib.py", line 630, in export_model
    export_archive.write_out(filepath)
  File "/home/pwuertz/.local/lib/python3.11/site-packages/keras/src/export/export_lib.py", line 474, in write_out
    self._filter_and_track_resources()
  File "/home/pwuertz/.local/lib/python3.11/site-packages/keras/src/export/export_lib.py", line 521, in _filter_and_track_resources
    tvs, ntvs = _list_variables_used_by_fns(fns)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pwuertz/.local/lib/python3.11/site-packages/keras/src/export/export_lib.py", line 809, in _list_variables_used_by_fns
    concrete_functions = [fn.get_concrete_function()]
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pwuertz/.local/lib/python3.11/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 1251, in get_concrete_function
    concrete = self._get_concrete_function_garbage_collected(*args, **kwargs)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pwuertz/.local/lib/python3.11/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 1221, in _get_concrete_function_garbage_collected
    self._initialize(args, kwargs, add_initializers_to=initializers)
  File "/home/pwuertz/.local/lib/python3.11/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 696, in _initialize
    self._concrete_variable_creation_fn = tracing_compilation.trace_function(
                                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pwuertz/.local/lib/python3.11/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compilation.py", line 178, in trace_function
    concrete_function = _maybe_define_function(
                        ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pwuertz/.local/lib/python3.11/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compilation.py", line 283, in _maybe_define_function
    concrete_function = _create_concrete_function(
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pwuertz/.local/lib/python3.11/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compilation.py", line 310, in _create_concrete_function
    traced_func_graph = func_graph_module.func_graph_from_py_func(
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pwuertz/.local/lib/python3.11/site-packages/tensorflow/python/framework/func_graph.py", line 1059, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pwuertz/.local/lib/python3.11/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 599, in wrapped_fn
    out = weak_wrapped_fn().__wrapped__(*args, **kwds)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pwuertz/.local/lib/python3.11/site-packages/tensorflow/python/eager/polymorphic_function/autograph_util.py", line 52, in autograph_handler
    raise e.ag_error_metadata.to_exception(e)
  File "/home/pwuertz/.local/lib/python3.11/site-packages/tensorflow/python/eager/polymorphic_function/autograph_util.py", line 41, in autograph_handler
    return api.converted_call(
           ^^^^^^^^^^^^^^^^^^^
  File "/home/pwuertz/.local/lib/python3.11/site-packages/tensorflow/python/autograph/impl/api.py", line 439, in converted_call
    result = converted_f(*effective_args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmp/__autograph_generated_file126962ye.py", line 44, in tf__error_handler
    ag__.if_stmt(ag__.not_(ag__.converted_call(ag__.ld(is_traceback_filtering_enabled), (), None, fscope)), if_body, else_body, get_state, set_state, ('do_return', 'retval_'), 2)
  File "/home/pwuertz/.local/lib/python3.11/site-packages/tensorflow/python/autograph/operators/control_flow.py", line 1217, in if_stmt
    _py_if_stmt(cond, body, orelse)
  File "/home/pwuertz/.local/lib/python3.11/site-packages/tensorflow/python/autograph/operators/control_flow.py", line 1270, in _py_if_stmt
    return body() if cond else orelse()
                               ^^^^^^^^
  File "/tmp/__autograph_generated_file126962ye.py", line 40, in else_body
    raise ag__.converted_call(ag__.ld(e).with_traceback, (ag__.ld(filtered_tb),), None, fscope) from None
  File "/tmp/__autograph_generated_file126962ye.py", line 34, in else_body
    retval_ = ag__.converted_call(ag__.ld(fn), tuple(ag__.ld(args)), dict(**ag__.ld(kwargs)), fscope)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: in user code:

    File "/home/pwuertz/.local/lib/python3.11/site-packages/keras/src/utils/traceback_utils.py", line 729, in error_handler  *
        filtered_tb = _process_traceback_frames(e.__traceback__)

    TypeError: Exception encountered when calling UpSampling2D.call().
    
    unsupported operand type(s) for *: 'NoneType' and 'int'
    
    Arguments received by UpSampling2D.call():
      • inputs=tf.Tensor(shape=(None, None, None, 1), dtype=float32)

The error is not triggered when using the default nearest UpSampling2D(), or when using fixed width/height inputs.

Building and training such a model still works, but I can't export it as TF-model or convert it via tf2onnx.

[fchollet]

Thanks for the report, this is fixed at HEAD.