Skip to content

Fix JAX implementation of Argmax #809

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 17 commits into from
Jul 6, 2024
Merged

Fix JAX implementation of Argmax #809

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
71fe3b9
Edited xfail reason for some 3 JAX tests
HangenYuu Jun 7, 2024
a170d33
Update tests/link/jax/test_extra_ops.py
HangenYuu Jun 17, 2024
d0fdd6e
Remove expected fail mark for test that actually passes
HangenYuu Jun 17, 2024
63b0832
Merge branch 'xfail_reason_update' of https://github.com/HangenYuu/py…
HangenYuu Jun 17, 2024
68e245f
Changed reason for xfail
HangenYuu Jun 28, 2024
3679791
Edited xfail reason for some 3 JAX tests
HangenYuu Jun 7, 2024
ae0801d
Remove expected fail mark for test that actually passes
HangenYuu Jun 17, 2024
f287d6c
Update tests/link/jax/test_extra_ops.py
HangenYuu Jun 17, 2024
8b05952
Changed reason for xfail
HangenYuu Jun 28, 2024
88b17b3
Merge branch 'xfail_reason_update' of github.com:HangenYuu/pytensor i…
HangenYuu Jul 3, 2024
349beef
Edited xfail reason for some 3 JAX tests
HangenYuu Jun 7, 2024
8ed6ab2
Remove expected fail mark for test that actually passes
HangenYuu Jun 17, 2024
b50a906
Update tests/link/jax/test_extra_ops.py
HangenYuu Jun 17, 2024
89c7bff
Changed reason for xfail
HangenYuu Jun 28, 2024
2c4a378
Rebase to remove MaxAndArgmax
HangenYuu Jul 3, 2024
c3c5fd0
Merge branch 'xfail_reason_update' of github.com:HangenYuu/pytensor i…
HangenYuu Jul 3, 2024
be9c9f8
Edited implementation of argmax. Changed test name for multiomni_output
HangenYuu Jul 6, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion tests/link/jax/test_extra_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ def test_extra_ops():

@pytest.mark.xfail(
version_parse(jax.__version__) >= version_parse("0.2.12"),
reason="Omnistaging cannot be disabled",
reason="JAX Numpy API does not support dynamic shapes",
)
def test_extra_ops_omni():
a = matrix("a")
Expand Down
4 changes: 2 additions & 2 deletions tests/link/jax/test_nlinalg.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,7 @@ def assert_fn(x, y):

@pytest.mark.xfail(
version_parse(jax.__version__) >= version_parse("0.2.12"),
reason="Omnistaging cannot be disabled",
reason="JAX Numpy API does not support dynamic shapes",
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Contributor Author

@HangenYuu HangenYuu Jun 17, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Member

@ricardoV94 ricardoV94 Jun 17, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)."

)
def test_jax_basic_multiout_omni():
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

image

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

# Test that a single output of a multi-output `Op` can be used as input to
Expand All @@ -96,7 +96,7 @@ def test_jax_basic_multiout_omni():

@pytest.mark.xfail(
version_parse(jax.__version__) >= version_parse("0.2.12"),
reason="Omnistaging cannot be disabled",
reason="`dot` -> `Gemv` optimization is incompatible with JAX",
)
def test_tensor_basics():
y = vector("y")
Expand Down
Loading