Fix JAX implementation of Argmax

[HangenYuu]

Description

"Omnistaging can't be disabled" was added as reason for some expected failures due to breaking changes introduced in JAX 0.2.12 onwards. As omnistaging already became default in JAX for 4 years, more descriptive reasons are needed for the tests affected (3 in total).

Related Issue

• Closes Update xfail reason "Omnistaging can't be disabled" in JAX tests #780
• Related to #

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 81.01%. Comparing base (5c8afae) to head (be9c9f8).
Report is 155 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #809      +/-   ##
==========================================
+ Coverage   80.97%   81.01%   +0.03%     
==========================================
  Files         169      170       +1     
  Lines       47015    46924      -91     
  Branches    11497    11497              
==========================================
- Hits        38072    38016      -56     
+ Misses       6728     6700      -28     
+ Partials     2215     2208       -7     

Files with missing lines	Coverage Δ
pytensor/link/jax/dispatch/nlinalg.py	97.77% <100.00%> (+6.76%)	⬆️

... and 39 files with indirect coverage changes

[ricardoV94]

Why does this test fail? Also seems like it should just work?

[HangenYuu]

I was looking at it. It actually failed, but not due to dynamic shape error with omnistaging. For this one, when I removed the xfail to see the underlying error, I got jax.errors.TracerIntegerConversionError: The __index__() method was called on traced array with shape int64[]. instead. Checking here I see that the error is due to passing in a traced value in place of a Python integer. To get rid of it, we need to mark the value at the function level. That is another PR I guess.

[HangenYuu]

Another problem I saw when I accidentally updated NumPy to 2.0 is many core APIs broke down. One example is np.obj2sctype(dtype) used in pytensor/tensor/type.py:105, which was removed in NumPy 2.0. This one actually made the test failed at collecting state. I raised an issue for this one.

[ricardoV94]

Hmm can you try to change x = dvector(shape=(2,)) and see if it works? Since that's not the concern of the test it's a shame to have it xfail

[ricardoV94]

PyTensor is explicitly not compatible with numpy 2.0 at the moment: #688

[HangenYuu]

dvector() has no shape argument so I tried vector(dtype="float64", shape=(2,)) instead. The error persists. The stack trace pointed the error to start from this temporary function generated when we turn the Op to JAX:

def jax_funcified_fgraph(tensor_variable):
    # MaxAndArgmax{axis=(0,)}(<Vector(float64, shape=(2,))>)
    tensor_variable_1, argmax = maxandargmax(tensor_variable)
    # Mul(max, argmax)
    tensor_variable_2 = elemwise_fn(tensor_variable_1, argmax)
    return (tensor_variable_2,)

The solution from here suggests that we need to somehow mark tensor_variable in the function above as static with functools.partial and jax.jit, but I think that it is out of scope of this PR? Maybe we can mark this as xfail for "Operation leads to JAX Tracer object is used in a context where a Python integer is expected (jax.errors.TracerIntegerConversionError)."

[ricardoV94]

Can you rebase, we got rid of MaxAndArgmax recently

[HangenYuu]

@ricardoV94 Sadly the jax.errors.TracerIntegerConversionError: The index() method was called on traced array with shape int64[] error persists 🥲.

[ricardoV94]

Okay so the source of this error is just the dispatch implementation of Argmax using Jax arrays for the transpose axis and the reshape operation. Changing those to use numpy arrays and converting to tuple before they are actually used allows the test to pass. Something like this:

[ricardoV94]

We should rename the test to be just test max_and_argmax, it is no longer a multioutput Operation, so the old name doesn't make sense. Also can you confirm this is the only test we had for Max / Argmax in JAX?

[HangenYuu]

Yes that's the only test for Max/Argmax we have.