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 16 commits
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
4 changes: 2 additions & 2 deletions 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,9 +65,9 @@ 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():
def test_extra_ops_dynamic_shapes():
a = matrix("a")
a.tag.test_value = np.arange(6, dtype=config.floatX).reshape((3, 2))

Expand Down
6 changes: 1 addition & 5 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="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 @@ -95,10 +95,6 @@ def test_jax_basic_multiout_omni():
compare_jax_and_py(out_fg, [np.r_[1, 2]])


@pytest.mark.xfail(
version_parse(jax.__version__) >= version_parse("0.2.12"),
reason="Omnistaging cannot be disabled",
)
def test_tensor_basics():
y = vector("y")
y.tag.test_value = np.r_[1.0, 2.0].astype(config.floatX)
Expand Down
Loading