Remove reliance on __jax_array__ to unwrap variables. (#21719)

JAX uses `__jax_array__` to handle non-JAX types. For instance when doing `a * v` where `a` is a `jax.Array` and `v` is a `keras.Variable`, the `jax.Array.__mul__` implementation calls `v.__jax_array__()` because `v` is not a JAX type.

However, `__jax_array__` did not work in all contexts, and the next version of JAX further restricts which contexts it works in.

The fix rarely involves explictly calling `v.value`. Instead, we rely on existing mechanisms that are already in place to unwrap variables in a lot of contexts:
- ops are always supposed to call `convert_to_tensor` on tensor inputs and `convert_to_tensor` extracts values from variables
- using `keras.ops` instead of native ops (+ - * / < > & etc.) unwraps variables. It is already a best practice to use `keras.ops` instead of native ops:
    - to support the creation of functional models via `KerasTensor`s and their serialization
    - to have consistent type promotion between backends
    - to support sparse tensors and ragged tensors

This was tested via a seperate PR #21702 that won't be submitted because of https://github.com/keras-team/keras/pull/21702/files#diff-900deadc65fc119ce93fb813e340dcb644b8eab9e7c0207bf37cdc05b8e8796e .

Author: hertschuh

keras/src/backend/jax/nn.py

@@ -396,6 +396,8 @@ def depthwise_conv(
     feature_group_count = (
         inputs.shape[-1] if data_format == "channels_last" else inputs.shape[1]
     )
+    kernel = convert_to_tensor(kernel)
+    inputs = convert_to_tensor(inputs)
     kernel = jnp.reshape(
         kernel,
         kernel.shape[:-2] + (1, feature_group_count * kernel.shape[-1]),

keras/src/backend/jax/numpy.py

@@ -543,15 +543,18 @@ def clip(x, x_min, x_max):
 
 def concatenate(xs, axis=0):
     bcoo_count = builtins.sum(isinstance(x, jax_sparse.BCOO) for x in xs)
-    if bcoo_count:
-        if bcoo_count == len(xs):
-            axis = canonicalize_axis(axis, len(xs[0].shape))
-            return jax_sparse.bcoo_concatenate(xs, dimension=axis)
-        else:
-            xs = [
-                x.todense() if isinstance(x, jax_sparse.JAXSparse) else x
-                for x in xs
-            ]
+    if bcoo_count == len(xs):
+        axis = canonicalize_axis(axis, len(xs[0].shape))
+        return jax_sparse.bcoo_concatenate(xs, dimension=axis)
+    elif bcoo_count:
+        xs = [
+            x.todense()
+            if isinstance(x, jax_sparse.JAXSparse)
+            else convert_to_tensor(x)
+            for x in xs
+        ]
+    else:
+        xs = [convert_to_tensor(x) for x in xs]
     return jnp.concatenate(xs, axis=axis)
 
 
@@ -1087,6 +1090,7 @@ def reshape(x, newshape):
         if None not in output_shape:
             newshape = output_shape
         return jax_sparse.bcoo_reshape(x, new_sizes=newshape)
+    x = convert_to_tensor(x)
     return jnp.reshape(x, newshape)
 
 
@@ -1149,10 +1153,12 @@ def sort(x, axis=-1):
 
 
 def split(x, indices_or_sections, axis=0):
+    x = convert_to_tensor(x)
     return jnp.split(x, indices_or_sections, axis=axis)
 
 
 def stack(x, axis=0):
+    x = [convert_to_tensor(t) for t in x]
     return jnp.stack(x, axis=axis)
 
 
@@ -1338,6 +1344,7 @@ def squeeze(x, axis=None):
             axis = tuple(i for i, d in enumerate(x.shape) if d == 1)
         axis = to_tuple_or_list(axis)
         return jax_sparse.bcoo_squeeze(x, dimensions=axis)
+    x = convert_to_tensor(x)
     return jnp.squeeze(x, axis=axis)
 
 

keras/src/backend/jax/optimizer.py

@@ -36,13 +36,14 @@ def _backend_apply_gradients(self, grads, trainable_variables):
             new_g_accs = jax.lax.cond(
                 is_update_step,
                 lambda: [jnp.zeros(g.shape, dtype=g.dtype) for g in acc_grads],
-                lambda: [g + acc_g for g, acc_g in zip(grads, acc_grads)],
+                lambda: [g + acc_g.value for g, acc_g in zip(grads, acc_grads)],
             )
 
             grads = jax.lax.cond(
                 is_update_step,
                 lambda: [
-                    (g + acc_g) / steps for g, acc_g in zip(grads, acc_grads)
+                    (g + acc_g.value) / steps
+                    for g, acc_g in zip(grads, acc_grads)
                 ],
                 lambda: list(grads),
             )

keras/src/layers/attention/attention.py

@@ -121,7 +121,7 @@ def _calculate_scores(self, query, key):
         if self.score_mode == "dot":
             scores = ops.matmul(query, ops.transpose(key, axes=[0, 2, 1]))
             if self.scale is not None:
-                scores *= self.scale
+                scores = ops.multiply(scores, self.scale)
         elif self.score_mode == "concat":
             # Reshape tensors to enable broadcasting.
             # Reshape into [batch_size, Tq, 1, dim].

keras/src/layers/layer_test.py

@@ -678,7 +678,7 @@ def __init__(self):
             def call(self, x):
                 # Should not autocast.
                 assertDType(self.v, "float32")
-                return ops.cast(x, "float32") + self.v
+                return ops.add(ops.cast(x, "float32"), self.v)
 
         # A layer that is explicitly full precision.
         class InnerLayerTwo(layers.Layer):
@@ -694,7 +694,7 @@ def __init__(self):
             def call(self, x):
                 # Should not autocast.
                 assertDType(self.v, "float32")
-                return x + self.v
+                return ops.add(x, self.v)
 
         # A layer that is explicitly mixed precision but with autocast=False
         # weight.
@@ -732,7 +732,7 @@ def call(self, x):
                 # Should autocast.
                 assertDType(self.v, "float16")
                 return self.inner_three(
-                    self.inner_two(self.inner_one(x + self.v))
+                    self.inner_two(self.inner_one(ops.add(x, self.v)))
                 )
 
         layer = MixedPrecisionLayer()
@@ -935,7 +935,7 @@ def call(self, x):
                 x = x + backend.random.normal(
                     shape=(), seed=self._seed_generator
                 )
-                return x + self.tw + self.ntw
+                return ops.add(x, ops.add(self.tw, self.ntw))
 
         data = np.random.random((3, 4))
         layer = TestLayer()

keras/src/layers/normalization/layer_normalization_test.py

@@ -99,8 +99,8 @@ def test_correctness(self):
         ).astype("float32")
 
         out = layer(inputs)
-        out -= layer.beta
-        out /= layer.gamma
+        out = ops.subtract(out, layer.beta)
+        out = ops.divide(out, layer.gamma)
 
         self.assertAllClose(ops.mean(out), 0.0, atol=1e-1)
         self.assertAllClose(ops.std(out), 1.0, atol=1e-1)

keras/src/layers/normalization/rms_normalization_test.py

@@ -38,10 +38,12 @@ def test_correctness(self):
         inputs = ops.convert_to_tensor(inputs)
 
         out = layer(inputs)
-        expected = (
-            inputs
-            * ops.rsqrt(ops.mean(ops.square(inputs), axis=-1, keepdims=True))
-            * layer.scale
+        expected = ops.multiply(
+            ops.multiply(
+                inputs,
+                ops.rsqrt(ops.mean(ops.square(inputs), axis=-1, keepdims=True)),
+            ),
+            layer.scale,
         )
 
         self.assertAllClose(out, expected, atol=1e-1)

keras/src/layers/rnn/gru.py

@@ -261,7 +261,7 @@ def call(self, inputs, states, training=False):
             matrix_x = ops.matmul(inputs, self.kernel)
             if self.use_bias:
                 # biases: bias_z_i, bias_r_i, bias_h_i
-                matrix_x += input_bias
+                matrix_x = ops.add(matrix_x, input_bias)
 
             x_z, x_r, x_h = ops.split(matrix_x, 3, axis=-1)
 

keras/src/layers/rnn/lstm.py

@@ -276,9 +276,9 @@ def call(self, inputs, states, training=False):
 
             z = ops.matmul(inputs, self.kernel)
 
-            z += ops.matmul(h_tm1, self.recurrent_kernel)
+            z = ops.add(z, ops.matmul(h_tm1, self.recurrent_kernel))
             if self.use_bias:
-                z += self.bias
+                z = ops.add(z, self.bias)
 
             z = ops.split(z, 4, axis=1)
             c, o = self._compute_carry_and_output_fused(z, c_tm1)

keras/src/layers/rnn/simple_rnn.py

@@ -160,7 +160,7 @@ def call(self, sequence, states, training=False):
             sequence = sequence * dp_mask
         h = ops.matmul(sequence, self.kernel)
         if self.bias is not None:
-            h += self.bias
+            h = ops.add(h, self.bias)
 
         if training and rec_dp_mask is not None:
             prev_output = prev_output * rec_dp_mask