-
Notifications
You must be signed in to change notification settings - Fork 19.4k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Unable to create model with Keras=3.2.1 and jax[cpu]=0.4.26 #19539
Comments
Update: Unfortunately it's not working anymore on Colab as well. Edit: Works properly with PyTorch and TensorFlow backend on Colab and local Windows machine. |
I've simplified the Colab code further and added print statements for debugging. The code fails all the time on the
|
Hi @veb-101, the 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 |
@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. |
Hello, I'm attempting to create the MobileViT model using the Keras 3 backend.
There are two main aspects:
Here is my code: Keras3_jax_windows.py
The code functions correctly on Colab (where I installed the latest versions of Jax and Keras), so I assume it should also work well on Linux machines.
On the Windows machine, it operates as intended with the TensorFlow backend. However, when using the Jax backend, the model creation fails. Interestingly, initializing the custom layer works properly, and I can even perform a forward pass through it.
Code snippet:
Output:
I know the code fails at this point:
The text was updated successfully, but these errors were encountered: