Following is the main query. Any suggestion on that would be appreciated. Thanks.

  

      github.com/keras-team/keras
  

  
    
  
	  
  

  
    
      InvalidArgumentError: 9 root error(s) found. 9
    

    
      
        opened 
        
          
        
        Nov 3, 2021
      
      


      
        
          
          innat
        
      
    

    
    
  


  
    https://discuss.tensorflow.org/tag/keras

**System information**.
- Have I wr…itten custom code (as opposed to using a stock example script provided in Keras):
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Colab
- TensorFlow installed from (source or binary): Colab
- TensorFlow version (use command below): 2.6
- Python version: Colab
- Bazel version (if compiling from source):
- GPU model and memory: TPU
- Exact command to reproduce:

You can collect some of this information using our environment capture script:

https://github.com/tensorflow/tensorflow/tree/master/tools/tf_env_collect.sh

**Describe the problem**.

Describe the problem clearly here. Be sure to convey here why it's a bug in Keras or why the requested feature is needed.

**Describe the current behavior**.

I was trying to run [`keras/vit`](https://keras.io/examples/vision/image_classification_with_vision_transformer/) code example on TPU but it causes the following error. 

> InvalidArgumentError: 9 root error(s) found.
  (0) Invalid argument: {{function_node __inference_train_function_49310}} Incompatible shapes: [256,25,64] vs. [144,64]

>	 [[{{node gradient_tape/model_1/patch_encoder_1/add/BroadcastGradientArgs}}]]
	 [[TPUReplicate/_compile/_12894506912996439067/_4]]
	 [[tpu_compile_succeeded_assert/_14236840403296971761/_5/_197]]
  (1) Invalid argument: {{function_node __inference_train_function_49310}} Incompatible shapes: [256,25,64] vs. [144,64]

>	 [[{{node gradient_tape/model_1/patch_encoder_1/add/BroadcastGradientArgs}}]]
	 [[TPUReplicate/_compile/_12894506912996439067/_4]]
	 [[tpu_compile_succeeded_assert/_14236840403296971761/_5/_245]]
  (2) Invalid argument: {{function_node __inference_train_function_49310}} Incompatible shapes: [256,25,64] vs. [144,64]


 I think it's due to the use of `tf.image.extract_patches` function on TPU. But not sure. 

## [Reproducible Coe. ](https://colab.research.google.com/drive/1DAQpG9YD-5vH-A-jfaISwFzMkE6oAlm9?usp=sharing)

**Describe the expected behavior**.

It should simply run. 

**[Contributing](https://github.com/keras-team/keras/blob/master/CONTRIBUTING.md)**.

- Do you want to contribute a PR? (yes/no):
- If yes, please read [this page](https://github.com/keras-team/keras/blob/master/CONTRIBUTING.md) for instructions
- Briefly describe your candidate solution(if contributing):

**Standalone code to reproduce the issue**.

Provide a reproducible test case that is the bare minimum necessary to generate
the problem. If possible, please share a link to Colab/Jupyter/any notebook.


**Source code / logs**.

Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached. Try to provide a reproducible test case that is the bare minimum necessary to generate the problem.