EHN: cluster: JAX support (non-jitted)

[crusaderky]

• Part of ENH: tracking issue for JAX support #22246
• See also ENH: special: JAX support (non-jitted) #22256

[crusaderky]

This fails on jax.jit. My current intention is to change jax.jit itself to special-case arr.at[idx].set(value) when idx is a boolean mask and value is a scalar, so that it can be rewritten as jnp.where(idx, value, arr). Failing that, I can implement the same special case in array-api-extra.

[jakevdp]

My current intention is to change jax.jit itself

I don't think changing jax.jit itself is a viable path here – I wouldn't suggest starting on that route.

It might be viable to make arr.at[idx].set(value) lower to lax.select rather than lax.scatter in the specific case of a boolean idx. I've tried that in the past, but it's really tricky to properly handle all corner cases of broadcasted and/or multi-dimensional indices, correctly implementing autodiff and batching rules, etc.

The easiest thing would probably be to do this at the level of xp.at, though boolean indices were specifically excluded from the initial discussions there.

[crusaderky]

I've tried that in the past

Do you have a (partial, non-functioning) PR I could start from?

[crusaderky]

This fails on jax.jit with an error message that's not not useful to final scipy users, unless you explicitly pass check_finite=False.
data-apis/array-api-compat#225 is a blocker for a more robust fix.

[crusaderky]

All CI failures are unrelated

[lucascolley]

thanks @crusaderky, LGTM!

Are there any potential improvements to make to

scipy/scipy/cluster/hierarchy.py

Lines 1839 to 1843 in 7a81451

	if is_jax(xp):
	# calculate_cluster_sizes doesn't accept read-only arrays
	Zpart = np.array(Zpart, copy=True)
	else:
	Zpart = np.asarray(Zpart)

at this time?

[lucascolley]

also, have you ran the full CUDA tests locally? I can, if not.

[crusaderky]

thanks @crusaderky, LGTM!

Are there any potential improvements to make to

scipy/scipy/cluster/hierarchy.py

Lines 1839 to 1843 in 7a81451

	if is_jax(xp):
	# calculate_cluster_sizes doesn't accept read-only arrays
	Zpart = np.array(Zpart, copy=True)
	else:
	Zpart = np.asarray(Zpart)

at this time?

Not yet. I intend to replace this and all other calls to np.asarray in the middle of xp-aware code with xpx.apply_numpy_function(func, arg), which internally will call jax.pure_callback(func, arg) for JAX and xp.asarray(func(np.asarray(arg)) for everything else.

[crusaderky]

also, have you ran the full CUDA tests locally? I can, if not.

Yes. All green functionally. This said, there are points where arrays are yanked from device memory to host memory by np.asarray, ran through a Cython kernel (search for _vq.), and sent back to jax.

[lucascolley]

This said, there are points where arrays are yanked from device memory to host memory by np.asarray, ran through a Cython kernel (search for _vq.), and sent back to jax.

Ah, does JAX not raise/warn there? Do we want to error out there @rgommers as per gh-18286?

[jakevdp]

Ah, does JAX not raise/warn there?

JAX will raise there during traced execution (e.g. under JIT), but in eager execution np.asarray is just fine. If you want to prevent this, you can enable the transfer guard.

[rgommers]

Using the "disallow" transfer guard in scipy/conftest.py would seem reasonable to me.

[crusaderky]

Using the "disallow" transfer guard in scipy/conftest.py would seem reasonable to me.

I'll investigate this but I expect it to result in a substantial amount of xfails. It would also make sense to figure out the equivalent setting for torch and cupy within the same scope. Either way, I'd rather leave it to a later PR...

[lucascolley]

I'll investigate this but I expect it to result in a substantial amount of xfails. It would also make sense to figure out the equivalent setting for torch and cupy within the same scope. Either way, I'd rather leave it to a later PR...

It should line up with existing cpu_only skips, no? CuPy errors on np.asarray

[rgommers]

It should line up with existing cpu_only skips, no? CuPy errors on np.asarray

That's what I expect as well. Same for PyTorch, it does not do implicit transfers.

[lucascolley]

Using the "disallow" transfer guard in scipy/conftest.py would seem reasonable to me.

Reminder to follow-up with this - thanks!