Added support for the jax.numpy.argsort

PennyLaneAI / catalyst

A JIT compiler for hybrid quantum programs in PennyLane

https://docs.pennylane.ai/projects/catalyst

Apache License 2.0

101 stars 26 forks source link

Added support for the jax.numpy.argsort #901

Closed mehrdad2m closed 3 days ago

mehrdad2m commented 3 days ago

Context:

jax.numpy.argsort is currently not supported in Catalyst, therefore compiling a function with @qjit that includes argsort would result in an error, see issue #810

Description of the Change:

Added hlo-legalize-sort to the HLO_LOWERING_PASS pipeline.
Added a pytest to confirm the functionality of jax.numpy.argsort
Removed jax.numpy.argsort from the unsopported jax functions

Fixes #810

erick-xanadu commented 3 days ago

@mehrdad2m awesome. It looks like there's a conflict in the changelog and can you please add context and a description of the change above? The PR description will be the commit message. After that happy to approve :)

mehrdad2m commented 3 days ago

@mehrdad2m awesome. It looks like there's a conflict in the changelog and can you please add context and a description of the change above? The PR description will be the commit message. After that happy to approve :)

Done!