fix failing jax test

[divyashreepathihalli]

/~https://github.com/keras-team/keras-hub/actions/runs/12173106768/job/33952917443?pr=1998

Log dump :
=================================== FAILURES ===================================
____________________ CLIPBackboneTest.test_backbone_basics _____________________

self = <keras_hub.src.models.clip.clip_backbone_test.CLIPBackboneTest testMethod=test_backbone_basics>

def test_backbone_basics(self):

  self.run_backbone_test(

        cls=CLIPBackbone,
        init_kwargs=self.init_kwargs,
        input_data=self.input_data,
        expected_output_shape={
            "vision_logits": (2, 2),
            "text_logits": (2, 2),
        },
    )

keras_hub/src/models/clip/clip_backbone_test.py:29:

---

keras_hub/src/tests/test_case.py:444: in run_backbone_test
backbone = cls(**init_kwargs)
keras_hub/src/models/clip/clip_backbone.py:207: in init
vision_logits, text_logits = self.clip_head(
/opt/hostedtoolcache/Python/3.9.20/x64/lib/python3.9/site-packages/keras/src/utils/traceback_utils.py:113: in error_handler
return fn(*args, **kwargs)
/opt/hostedtoolcache/Python/3.9.20/x64/lib/python3.9/site-packages/keras/src/layers/layer.py:826: in call
self._maybe_build(call_spec)
/opt/hostedtoolcache/Python/3.9.20/x64/lib/python3.9/site-packages/keras/src/layers/layer.py:1368: in _maybe_build
self.build(**shapes_dict)
/opt/hostedtoolcache/Python/3.9.20/x64/lib/python3.9/site-packages/keras/src/layers/layer.py:226: in build_wrapper
original_build_method(*args, **kwargs)
keras_hub/src/models/clip/clip_backbone.py:67: in build
self.logit_scale = self.add_weight(
/opt/hostedtoolcache/Python/3.9.20/x64/lib/python3.9/site-packages/keras/src/layers/layer.py:541: in add_weight
variable = backend.Variable(
/opt/hostedtoolcache/Python/3.9.20/x64/lib/python3.9/site-packages/keras/src/backend/common/variables.py:175: in init
self._initialize_with_initializer(initializer)
/opt/hostedtoolcache/Python/3.9.20/x64/lib/python3.9/site-packages/keras/src/backend/common/variables.py:373: in _initialize_with_initializer
self._initialize(value)

---

self =
value = 2.659260036932778

def _initialize(self, value):
    # Note that variable.shape is needed by distribution_lib

  self._shape = tuple(value.shape)

E AttributeError: 'float' object has no attribute 'shape'

/opt/hostedtoolcache/Python/3.9.20/x64/lib/python3.9/site-packages/keras/src/backend/jax/core.py:24: AttributeError
=========================== short test summary info ============================
FAILED keras_hub/src/models/clip/clip_backbone_test.py::CLIPBackboneTest::test_backbone_basics - AttributeError: 'float' object has no attribute 'shape'
=========== 1 failed, 968 passed, 832 skipped in 1481.21s (0:24:41) ============

[james77777778]

The commit: keras-team/keras@aae0520
causes the CI to fail.

Will submit a fix soon in Keras