keras.ops.shape doesn't support input with unknown dimension

[breadbread1984]

keras.ops.shape doesn't support input with unknown dimension with torch and tensorflow backend. jax backend is ok. the following code can reproduce the problem. please support input with unknown dimension for keras.ops.shape.

from os import environ
#environ['KERAS_BACKEND'] = 'tensorflow' # FAIL
#environ['KERAS_BACKEND'] = 'torch' # FAIL
environ['KERAS_BACKEND'] = 'jax' # OK
import keras as k

def Test():
   inputs = k.Input((None, 3))
   shape = k.layers.Lambda(lambda x: k.ops.shape(x))(inputs)
   return k.Model(inputs = inputs, outputs = shape)

import numpy as np
inputs = np.random.normal(size = (1,10,3))
outputs = Test()(inputs)
print(outputs)

the pure tensorflow implement has no such problem.

import tensorflow as tf

def Test():
   inputs = tf.keras.Input((None, 3))
   shape = tf.keras.layers.Lambda(lambda x: tf.shape(x))(inputs)
   return tf.keras.Model(inputs = inputs, outputs = shape)

import numpy as np
inputs = np.random.normal(size = (1,10,3))
outputs = Test()(inputs)
print(outputs)

[AakashKumarNain]

Can you paste the full trace here?

[breadbread1984]

tensorflow:

Traceback (most recent call last):
  File "/home/xieyi/test.py", line 14, in <module>
    outputs = Test()(inputs)
  File "/home/xieyi/test.py", line 9, in Test
    shape = k.layers.Lambda(lambda x: k.ops.shape(x))(inputs)
  File "/home/xieyi/anaconda3/lib/python3.9/site-packages/keras/src/utils/traceback_utils.py", line 123, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/home/xieyi/anaconda3/lib/python3.9/site-packages/keras/src/layers/core/lambda_layer.py", line 97, in compute_output_shape
    raise NotImplementedError(
NotImplementedError: Exception encountered when calling Lambda.call().

torch:

Traceback (most recent call last):
  File "/home/xieyi/test.py", line 14, in <module>
    outputs = Test()(inputs)
  File "/home/xieyi/test.py", line 9, in Test
    shape = k.layers.Lambda(lambda x: k.ops.shape(x))(inputs)
  File "/home/xieyi/anaconda3/lib/python3.9/site-packages/keras/src/utils/traceback_utils.py", line 123, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/home/xieyi/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/xieyi/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
NotImplementedError: Exception encountered when calling Lambda.call().

[SuryanarayanaY]

Hi @breadbread1984 ,

I have replicated the reported behaviour and attached the gist here for reference. With jax backend the code works fine but with torch and tensorflow raising the error as reported

[hertschuh]

@breadbread1984 ,

Thank you for reporting the issue.

The core of the issue is that keras.ops.shape is not a normal op in the sense that it doesn't return a tensor. Instead, it returns a tuple containing ints for known dimensions. keras.ops.shape(x) is just a way to write x.shape in a way that works across backends. In that sense keras.ops.shape(x) and tf.shape(x) are really not the same, which I agree is very confusing.

One easy way to fix the issue is to add convert_to_tensor (this fixes both TensorFlow and Torch):

shape = k.layers.Lambda(lambda x: k.backend.convert_to_tensor(k.ops.shape(x)))(inputs)

As far as TensorFlow is concerned, it's actually not the unknown dimensions that are the issue, the issue is the known dimensions (because of the way that unknown dimensions are supported). If you change the input to k.Input((None, None)), it actually works with TensorFlow (not Torch).

def Test():
   inputs = k.Input((None, None))
   shape = k.layers.Lambda(lambda x: k.ops.shape(x))(inputs)
   return k.Model(inputs = inputs, outputs = shape)

import numpy as np
inputs = np.random.normal(size = (1,10,3))
outputs = Test()(inputs)

That being said, I have a fix for TensorFlow that should make your code work as-is. I will report back.

[breadbread1984]

OK thx for the solution

[pwuertz]

That being said, I have a fix for TensorFlow that should make your code work as-is. I will report back.

Any update on this? I also have a dynamically sized model with input-shape (None, None, 3) that requires tf.shape ops at some point. Migrating this model from TF/Keras2 to Keras3 still seems to be an issue.

Keras 3.3.2 doesn't handle the aforementioned construct:

inputs = k.Input(shape=(None, None, 3), dtype=np.float32)
shape = k.layers.Lambda(lambda x: k.ops.shape(x))(inputs)

# raises NotImplementedError:
#   "We could not automatically infer the shape of the Lambda's output.
#   Please specify the `output_shape` argument for this Lambda layer."

The k.backend.convert_to_tensor solution suggested by @hertschuh also raises the same error.

[hertschuh]

@pwuertz ,

I tested with Keras 3.3.2 and this works with the Tensorflow backend and the JAX backend:

import keras as k
import numpy as np

def Test():
     inputs = k.Input(shape=(None, None, 3), dtype=np.float32)
     shape = k.layers.Lambda(lambda x: k.ops.convert_to_tensor(k.ops.shape(x)))(inputs)
     return k.Model(inputs = inputs, outputs = shape)

inputs = np.random.normal(size=(1, 2, 10, 3))
outputs = Test()(inputs)

In this case, you do need the convert_to_tensor. Can you confirm this also works for you?

[pwuertz]

Oh right, this indeed works with Tensorflow. I guess my notebook jumped over to the torch backend earlier, which still raises the NotImplementedError.

... which is still somewhat of an issue for me because Keras 2 to 3 also breaks tf2onnx.convert.from_keras and taking the multi-framework route using torch.onnx.export as an alternative doesn't work due to this bug / missing feature.

I think what we're really asking for here is some sort of k.layers.Shape operator akin to tf.shape or ONNX Shape?

[hertschuh]

@pwuertz ,
A changed was just pushed that fixes the code above for Torch. Can you confirm this works for your ONNX export?

I think what we're really asking for here is some sort of k.layers.Shape operator akin to tf.shape or ONNX Shape?

Hopefully this works now: k.layers.Lambda(lambda x: k.ops.convert_to_tensor(k.ops.shape(x)))

[pwuertz]

@pwuertz , A changed was just pushed that fixes the code above for Torch. Can you confirm this works for your ONNX export?

Confirmed, this now works in both TF and Torch, thanks!