Fix PARSeq decoder for TFLite compatibility (#2467)

* Fix PARSeq decoder for TFLite compatibility

Use ops.cond() instead of Python if-statement in PARSeqDecoder.call()
to ensure graph mode compatibility during TFLite conversion.

- Changed 'if tokens_length > 1:' to ops.cond(tokens_length > 1, ...)
- This allows TensorFlow to properly trace the graph with symbolic tensors

This fixes the 'Using a symbolic tf.Tensor as a Python bool is not allowed'
error during TFLite conversion when sequence length is dynamic (None).

* Fix PARSeq decoder for TFLite and JAX compatibility

* Fix PARSeq decoder for TFLite and JAX compatibility

* Update parseq_decoder.py

* Update parseq_decoder.py

* Refactor content and query embedding logic in PARSeqDecoder

Simplifies content and query embedding construction for better compatibility with JAX/TF graph backends. Removes dynamic slicing and Python conditionals, using ops.take and shape-based indexing to ensure consistent tensor shapes.

Author: pctablet505

keras_hub/src/models/parseq/parseq_decoder.py

@@ -364,20 +364,32 @@ def call(
         null_context = self.hidden_dim**0.5 * self.token_embedding(
             token_ids[:, :1]
         )
-        if tokens_length > 1:
-            content = self.pos_query_embeddings[:, : tokens_length - 1, :]
-            content = content + self.hidden_dim**0.5 * self.token_embedding(
-                token_ids[:, 1:]
-            )
-            content = ops.concatenate([null_context, content], axis=1)
-        else:
-            content = null_context
+
+        # Build content embeddings. When tokens_length == 1, this produces an
+        # empty tensor (shape: bs, 0, hidden), avoiding the need for a Python
+        # conditional.
+        content_embeddings = self.hidden_dim**0.5 * self.token_embedding(
+            token_ids[:, 1:]
+        )
+        # Use ops.take instead of dynamic slicing for JAX/TF graph compatibility
+        pos_embeds = ops.take(
+            self.pos_query_embeddings,
+            ops.arange(ops.shape(content_embeddings)[1], dtype="int32"),
+            axis=1,
+        )
+        content = ops.concatenate(
+            [null_context, pos_embeds + content_embeddings], axis=1
+        )
 
         content = self.dropout(content)
 
         query = ops.multiply(
             ops.ones((bs, 1, 1), dtype=self.dtype),
-            self.pos_query_embeddings[:, :tokens_length, :],
+            ops.take(
+                self.pos_query_embeddings,
+                ops.arange(tokens_length, dtype="int32"),
+                axis=1,
+            ),
         )
         query = self.dropout(query)