Replace deprecated jax.core.ConcreteArray

[uriballo]

Context: Replace the deprecated jax.core.ConcreteArray from the is_abstract function in utils.py by a check with jax.numpy.

Description of the Change: APIs in jax.core are not considered public, and may be changed in future releases without warning (see the API compatibility policy) -> Issue from jax repo

Benefits: We no longer use jax.core private API.

Possible Drawbacks: Maybe could be made more efficient.

Related GitHub Issues: N/A

[josh146]

Hi @uriballo! Thanks for making this PR. For context, are you trying to use PL with a newer version of JAX beyond 0.4.28?

Currently, we are supporting JAX 0.4.28 with both Catalyst and PennyLane, but making plans to upgrade our support

[uriballo]

Hi @josh146, yes, I am trying to use PL with a newer version of JAX. However, as far as I know jax.core.ConcreteArray was never meant to be used publicly, it was just an internal implementation that got removed in version 0.4.36.

[josh146]

Got it, thanks @uriballo! Due to rapid development of PennyLane and Catalyst, at the moment we are testing and ensuring PennyLane works best with JAX 0.4.28, so for the best experience I recommend downgrading your local jax and jaxlib installations.

At the same time though we're looking at how to best upgrade the stack to work with the most recent JAX versions -- meanwhile @mudit2812 is likely the best person to help on this PR in particular.

[mudit2812]

Hey @uriballo! Thanks for working on this. Unfortunately, uses of JAX internals (which are now broken with jax > 0.4.28) are present in many other submodules. Because of this, enabling compatibility with jax > 0.4.28 will require a significantly bigger effort in the PennyLane codebase and also the rest of the ecosystem. This is something on our radar and we're currently scoping the changes and effort needed for this. For now, I can only recommend using JAX 0.4.28 when using pennylane as suggested by @josh146.

May I ask if anything is blocking you from downgrading JAX? Any information you may have about PennyLane/JAX incompatibilities can be helpful for us to track down changes needed on our end.

[uriballo]

Hi @mudit2812! Thanks for your reply. I haven't had any issues when using newer versions of JAX, other than the deprecated error whenever I tried to use JIT and I was not specifiying jax-jit as the interface. I'll follow your suggestions and downgrade to the last supported version. Thanks for your time and hopefully I'll be able to help when you decide to upgrade to newer JAX version :)