Remove reliance on __jax_array__ to unwrap variables.

[hertschuh]

Several months ago JAX deprecated passing __jax_array__-implementing objects directly, to, e.g., a jit-ted function. JAX has emitted a warning since that time. In a future release of JAX this will become a hard error.

[gemini-code-assist]

Summary of Changes

Hello @hertschuh, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses an upcoming breaking change in JAX by removing the direct use of the deprecated __jax_array__ method. It refactors several parts of the Keras JAX backend, layers, and optimizers to use explicit ops functions for arithmetic operations instead of Python's overloaded operators. Additionally, it enhances tensor conversion consistency in JAX backend functions and refines how optimizer variables are handled, ensuring compatibility and robustness with future JAX releases.

Highlights

• JAX __jax_array__ Deprecation: Removed direct reliance on the __jax_array__ method, which JAX has deprecated and will soon become a hard error. The method now explicitly raises a ValueError to prevent its use, signaling that alternative unwrapping mechanisms should be employed.
• Explicit ops Usage: Replaced implicit Python arithmetic operators (e.g., +, *, -, /) with explicit ops functions (e.g., ops.add, ops.multiply, ops.subtract, ops.divide) across various Keras layers, optimizers, and tests. This change ensures consistent backend operation handling and avoids potential issues with operator overloading.
• Tensor Conversion Consistency: Introduced convert_to_tensor calls in JAX backend functions like concatenate, reshape, split, stack, and squeeze. This ensures that inputs are consistently JAX tensors, preventing potential issues when non-tensor inputs are provided.
• Optimizer Variable Handling Refinement: Updated JAX optimizer logic to correctly access the underlying value of accumulated gradients (acc_g.value) and to pass variable values (v.value) in stateless optimizer applications, improving robustness and compatibility.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request is a great initiative to remove the reliance on the deprecated __jax_array__ for unwrapping variables, aligning the codebase with future JAX releases. The changes are extensive, touching many files to make implicit conversions explicit, primarily by replacing standard Python operators with keras.ops functions and explicitly unwrapping Variable objects. This significantly improves code clarity and robustness.

I've identified a critical issue with some leftover debugging code that must be removed, and a high-severity bug in the implementation of the stack function. Apart from these points, the changes look solid and consistent with the goal of the pull request.

[gemini-code-assist]

This debugging code, including the TODO UNDO comment, print statements, and the ValueError exception, should be removed before merging the pull request. It appears to have been added to trace usages of __jax_array__ but is not intended for the final version.

Suggested change

	# TODO UNDO
	import traceback
	print("### __jax_array__")
	traceback.print_stack()
	raise ValueError("__jax_array__")
	return self.value

[codecov-commenter]

Codecov Report

❌ Patch coverage is 87.09677% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.59%. Comparing base (cc56474) to head (889d5bf).
⚠️ Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
keras/src/backend/jax/core.py	0.00%	4 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #21702      +/-   ##
==========================================
- Coverage   82.59%   82.59%   -0.01%     
==========================================
  Files         572      572              
  Lines       58314    58324      +10     
  Branches     9130     9130              
==========================================
+ Hits        48166    48171       +5     
- Misses       7817     7821       +4     
- Partials     2331     2332       +1     

Flag	Coverage Δ
keras	82.39% <87.09%> (-0.01%)	⬇️
keras-jax	63.30% <87.09%> (-0.01%)	⬇️
keras-numpy	57.64% <48.38%> (-0.01%)	⬇️
keras-openvino	34.32% <0.00%> (-0.01%)	⬇️
keras-tensorflow	64.03% <41.93%> (-0.02%)	⬇️
keras-torch	63.61% <48.38%> (-0.02%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.