Repository: keras-team/keras

Issue:
`keras.ops.broadcast_to` throws error `expected ArrayLike, got KerasVariable` when broadcasting a Keras layer weight. Only on `jax` backend.
Hello,
When I try to broadcast a Keras layer's weight that I create with the `self.add_weight` method, using the `keras.ops.broadcast_to` function, everything works fine on both `torch` and `tensorflow` backends, but on `jax` backend,  I run into the following error:
```
Traceback (most recent call last):
  File "/home/abaid/projects/sandbox/jax/brodcast_error_with_keras.py", line 22, in <module>
    outputs = tst_layer(inputs)
              ^^^^^^^^^^^^^^^^^
  File "/home/abaid/miniconda3/envs/ML/lib/python3.11/site-packages/keras/src/utils/traceback_utils.py", line 123, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/home/abaid/projects/sandbox/jax/brodcast_error_with_keras.py", line 15, in call
    x_weight_broadcasted = ops.broadcast_to(self.x_weight, (8, 2))
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/abaid/miniconda3/envs/ML/lib/python3.11/site-packages/jax/_src/numpy/lax_numpy.py", line 1227, in broadcast_to
    return util._broadcast_to(array, shape)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/abaid/miniconda3/envs/ML/lib/python3.11/site-packages/jax/_src/numpy/util.py", line 413, in _broadcast_to
    arr = arr if isinstance(arr, Array) else lax.asarray(arr)
                                             ^^^^^^^^^^^^^^^^
  File "/home/abaid/miniconda3/envs/ML/lib/python3.11/site-packages/jax/_src/lax/lax.py", line 137, in asarray
    raise TypeError(f"asarray: expected ArrayLike, got {x} of type {type(x)}.")
TypeError: Exception encountered when calling TestLayer.call().

asarray: expected ArrayLike, got <KerasVariable shape=(1, 2), dtype=float32, path=test_layer/variable> of type <class 'keras.src.backend.jax.core.Variable'>.

Arguments received by TestLayer.call():
  • inputs=jnp.ndarray(shape=(8, 5), dtype=float32)
```

Here is the simplified code snippet that replicates the issue and produces the error traceback posted above.

```
import os
os.environ["KERAS_BACKEND"] = "jax"
from keras import ops, random, layers


class TestLayer(layers.Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.x_weight = self.add_weight(
            shape=(1, 2),
            initializer="random_normal"
        )

    def call(self, inputs):
        x_weight_broadcasted = ops.broadcast_to(self.x_weight, (8, 2))
        outputs = ops.concatenate([x_weight_broadcasted, inputs], axis=1)
        return outputs


tst_layer = TestLayer()
inputs = random.normal((8, 5))
outputs = tst_layer(inputs)
assert ops.shape(outputs) == (8, 7)
```


Discussion:
If I explicitly call the `.numpy()` method of the weight when passing its value the `ops.broadcast_to` method, it works fine. But not sure if that is how it's supposed to be used.

```
x_weight_broadcasted = ops.broadcast_to(self.x_weight.numpy(), (8, 2)) 	# works fine for all backends
```
Hi @KhawajaAbaid ,

I have replicated the reported bug with Jax backend.Other backend works fine.

Attached [gist](https://colab.research.google.com/gist/SuryanarayanaY/0333554c4192f2d99f0f0fbbde0f97ed/19116.ipynb) for reference with probable fix. 

This Op with jax backend calls jax.numpy.broadcast_to which expects array like object as input but we are passing Keras Variable which is a Tensor.

I am proposing a probable fix for this.

Resolving patch:
diff --git a/keras/backend/jax/numpy.py b/keras/backend/jax/numpy.py
--- a/keras/backend/jax/numpy.py
+++ b/keras/backend/jax/numpy.py
@@ -322,6 +322,7 @@ def average(x, axis=None, weights=None):
 
 
 def broadcast_to(x, shape):
+    x = convert_to_tensor(x)
     return jnp.broadcast_to(x, shape)
 
 


Test patch: