Skip to content

Rewriting the kron function using JAX implementation #684

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 14 commits into from
Mar 27, 2024
Merged

Rewriting the kron function using JAX implementation #684

Changes from 1 commit
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
8875fd0
Moved the kron function as-is from slinalg.py to nlinalg.py
tanish1729 Mar 26, 2024
9eb50a6
Updated the kron function to use JAX implementation and moved tests t…
tanish1729 Mar 26, 2024
be52b59
Updated test for kron
tanish1729 Mar 26, 2024
52ed8c3
Updated slinalg to use the new kron implementation
tanish1729 Mar 26, 2024
1611f9a
running all pre commit checks
tanish1729 Mar 26, 2024
5aef097
removed importing reshape
tanish1729 Mar 26, 2024
5299b75
added test for commute with inv
tanish1729 Mar 26, 2024
3e0c79c
updated test
tanish1729 Mar 27, 2024
a7902fd
working test for commute_with_inv
tanish1729 Mar 27, 2024
e3b1aaa
working test for commute_with_inv
tanish1729 Mar 27, 2024
3ea1cdb
skipped some half precision tests
tanish1729 Mar 27, 2024
7fa03cf
skipped some half precision tests
tanish1729 Mar 27, 2024
5104494
updated comment
tanish1729 Mar 27, 2024
35996a5
Parametrized the tests
tanish1729 Mar 27, 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
8 changes: 3 additions & 5 deletions tests/tensor/test_nlinalg.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -606,12 +606,10 @@ def test_perform(self):
np.testing.assert_allclose(out, np_val)

def test_kron_commutes_with_inv(self):
Copy link
Member

Choose a reason for hiding this comment

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

Since no code is reused, why not parametrize?

for shp0, shp1 in zip(
[(2, 3), (2, 3), (2, 4, 3)], [(6, 7), (4, 3, 5), (4, 3, 5)]
for i, (shp0, shp1) in enumerate(
zip([(2, 3), (2, 3), (2, 4, 3)], [(6, 7), (4, 3, 5), (4, 3, 5)])
):
if (pytensor.config.floatX == "float32") & (
len(shp0) == 3 or len(shp1) == 3
):
if (pytensor.config.floatX == "float32") & (i == 2):
Copy link
Member

Choose a reason for hiding this comment

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

Just skip everything with float32?

pytest.skip("Kron and inv do not commute at half precision")
x = tensor(dtype="floatX", shape=(None,) * len(shp0))
a = np.asarray(self.rng.random(shp0)).astype(config.floatX)
Expand Down