Fixes compute_output_shape for PaliGemmaVitEncoder and Gemma3VisionEncoderBlock

[JyotinderSingh]

The logic for handling a Keras 3.1 edge case was missing a return statement, causing this code to run into the following error when running keras_hub/src/models/pali_gemma/pali_gemma_backbone_test.py:

self = <PaliGemmaVitEncoder name=image_encoder, built=True>
inputs_shape = [None, None, None]

    def compute_output_shape(self, inputs_shape):
        if inputs_shape is None:
            # Fix the compatibility issue with Keras 3.1 where
            # `compute_output_spec` fails to propagate `inputs_shape`
            # correctly, causing it to be `None`.
            inputs_shape = [None, None, None]
        return [
            inputs_shape[0],
>           (inputs_shape[1] // self.patch_size) ** 2,
            self.hidden_dim,
        ]
E       TypeError: unsupported operand type(s) for //: 'NoneType' and 'int'

Additionally, we replace the last dimension of the output shape with self.hidden_dim, since we know for sure that we will be projecting hidden_dim outputs to the next layer.

The failing test now passes with these changes:

[abheesht17]

Thanks!

Is the UT which was failing earlier passing now?

[JyotinderSingh]

Thanks!

Is the UT which was failing earlier passing now?

Yes, verified locally with latest keras-team/keras master commit, which was failing earlier.

Attached failing and updated test results to the PR description

[abheesht17]

Can you make the same change here too, please?

https://github.com/keras-team/keras-hub/blob/master/keras_hub/src/models/gemma3/gemma3_vision_encoder.py#L491

[JyotinderSingh]

Can you make the same change here too, please?

https://github.com/keras-team/keras-hub/blob/master/keras_hub/src/models/gemma3/gemma3_vision_encoder.py#L491

Done