Unable to create model with Keras=3.2.1 and jax[cpu]=0.4.26 about keras

[Star9daisy]

Hi @veb-101, the MobileViT_v1_Block you create in colab looks similar to an official example here:
https://keras.io/examples/vision/image_classification_with_vision_transformer/

The calculation process is almost the same. So I compare yours with those in it. In short, you could change the following lines to make your codes work:

# Comment out this line
# self.patch_size_h, self.patch_size_w = kops.cast(self.patch_size_h, dtype="int32"), kops.cast(self.patch_size_w, dtype="int32")

# Replace ops.floor_divide with //
# n_patches_h = ops.floor_divide(height, self.patch_size_h)
# n_patches_w = ops.floor_divide(width, self.patch_size_w)
n_patches_h = height // self.patch_size_h
n_patches_w = width // self.patch_size_w

[veb-101]

Update: Unfortunately it's not working anymore on Colab as well.

colab link

Edit: Works properly with PyTorch and TensorFlow backend on Colab and local Windows machine.

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No

[veb-101]

I've simplified the Colab code further and added print statements for debugging. The code fails all the time on the reshape(...) function, and if I skip reshaping, then the transpose fails with the same error.

Shapes must be 1D sequences of concrete values of integer type, got (None, 32, 32, 144).

[veb-101]

@Star9daisy Thanks, it works. I was able to complete my first pretrained model port to Keras 3 finally

Such a small solution, and I was doing something extreme.

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No