Awkward JAX backend deprecation

[ianna]

We have been discussing the JAX backend issues on and off. Here’s the link to the initial post #3349.

[ikrommyd]

Pasting this here from the jax deprecation warning PR: #3609 (comment)
ak.jax.register_and_check() is the first line of code someone has to run before using any jax features. It should definitely raise a warning.

[ianna]

Pasting this here from the jax deprecation warning PR: #3609 (comment) ak.jax.register_and_check() is the first line of code someone has to run before using any jax features. It should definitely raise a warning.

We’re not preventing the use of JAX with Awkward; we’re removing the high-level Array JAX backend. The high-level functions ak.to_jax and ak.from_jax are staying, and they do rely on JAX nplike and backend components, which are changing. That said, I don’t think the warning should go in ak.jax.register_and_check(), since it’s the first thing users run and their code will likely continue to work unchanged. Placing the warning there could cause unnecessary confusion. It makes more sense to surface it directly in the nplike and backend layers, where the actual changes occur—this keeps the messaging accurate, targeted, and preserves user trust in the high-level API.

[ikrommyd]

register_and_check does this:

awkward/src/awkward/jax.py

Lines 31 to 53 in 2ca5ccd

	def register_and_check():
	"""
	Register Awkward Array node types with JAX's tree mechanism.
	"""
	try:
	import jax # noqa: TID251
	# ak.from_buffers needs this
	jax.config.update("jax_enable_x64", True)
	except ModuleNotFoundError:
	raise ModuleNotFoundError(
	"""install the 'jax' package with:
	python3 -m pip install jax jaxlib
	or
	conda install -c conda-forge jax jaxlib
	"""
	) from None
	_register()

and call this _register function:

awkward/src/awkward/jax.py

Line 83 in 2ca5ccd

def _register():

which I assume will not exist because awkward._connect.jax will be removed correct?

[ianna]

register_and_check does this:

awkward/src/awkward/jax.py

Lines 31 to 53 in 2ca5ccd

	def register_and_check():
	"""
	Register Awkward Array node types with JAX's tree mechanism.
	"""
	try:
	import jax # noqa: TID251
	# ak.from_buffers needs this
	jax.config.update("jax_enable_x64", True)
	except ModuleNotFoundError:
	raise ModuleNotFoundError(
	"""install the 'jax' package with:
	python3 -m pip install jax jaxlib
	or
	conda install -c conda-forge jax jaxlib
	"""
	) from None
	_register()

and call this _register function:

awkward/src/awkward/jax.py

Line 83 in 2ca5ccd

def _register():

which I assume will not exist because awkward._connect.jax will be removed correct?

Yes — removing the JAX backend means all JAX-related code in Awkward will need to be refactored or reconsidered.
The awkward._connect.jax module is currently planned for removal, but depending on ongoing discussions, we may need to keep parts of it (for example, to support ak.to_jax and ak._fromjax).

These conversion functions will definitely be maintained, but their implementation might change to reduce coupling with the full backend.

[ianna]

@ikrommyd - Regarding your concerns about the warning — it clearly appears during the tests, ensuring users see it early and often.
Overall, the change looks good to me.