Keras 3 expand dims not fully backend-agnostic

[MathiesW]

The function keras.ops.expand_dims(x, axis=(0, -1)) does not work with Tensorflow backend if more than one axis to expand along is defined. It raises InvalidArgumentError: 'dim' must be a tensor with a single value [Op:ExpandDims]. I guess it internally calls the function tensorflow.expand_dims() which only supports a single axis at a time. JAX backend works fine, Pytorch I did not try.

[james77777778]

The newest version should support this functionality (in all backends), which was actually introduced by me 😃

You can refer to this colab

from keras import backend
from keras import ops

print(backend.backend())  # tensorflow
x = ops.ones([2, 3, 4, 5])
print(x.shape)  # (2, 3, 4, 5)
x = ops.expand_dims(x, [0, -1])
print(x.shape)  # (1, 2, 3, 4, 5, 1)

You can also find the test for this:
https://github.com/keras-team/keras/blob/master/keras/ops/numpy_test.py#L3332-L3361

[SuryanarayanaY]

Hi @MathiesW ,

Referring to @james77777778 comment could you please confirm woth what Keras version you have tested? Please try with latest version and still problem persists let us know with repro code snippet. Thanks!

[MathiesW]

Hi @SuryanarayanaY and @james77777778
I was using Keras 3.0.5. Just updated to 3.1.1 and now it works as intended. Thank you!