Add JAX support for SortOp

[HarshvirSandhu]

Description

Implement JAX conversion for SortOp

Related Issue

• Closes FEAT: JAX Conversion for the given SortOp #595

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works

Type of change

• Maintenance

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 80.82%. Comparing base (082081a) to head (1ef5238).
Report is 26 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #657      +/-   ##
==========================================
+ Coverage   80.80%   80.82%   +0.02%     
==========================================
  Files         162      162              
  Lines       46743    46820      +77     
  Branches    11419    11438      +19     
==========================================
+ Hits        37770    37844      +74     
+ Misses       6731     6725       -6     
- Partials     2242     2251       +9     

Files	Coverage Δ
pytensor/link/jax/dispatch/tensor_basic.py	92.37% <100.00%> (+0.40%)	⬆️

... and 24 files with indirect coverage changes

[ricardoV94]

Looks good, I left just two small modification suggestions for readability and more extensive testing

[ricardoV94]

Suggested change

	def sort(arr, *args):
	return jnp.sort(arr, *args)
	def sort(arr, axis):
	return jnp.sort(arr, axis=axis)

[ricardoV94]

Suggested change

	def test_sort():
	x = matrix("x")
	out = pytensor.tensor.sort(x)
	fgraph = FunctionGraph([x], [out])
	arr = np.array([[1.0, 4.0], [5.0, 2.0]])
	compare_jax_and_py(fgraph, [arr])
	@pytest.mark.parametrize("axis", [None, -1])
	def test_sort(axis):
	x = matrix("x", shape=(2, 2), dtype="float64")
	out = pytensor.tensor.sort(x, axis=axis)
	fgraph = FunctionGraph([x], [out])
	arr = np.array([[1.0, 4.0], [5.0, 2.0]])
	compare_jax_and_py(fgraph, [arr])

[ricardoV94]

Thanks @HarshvirSandhu