Skip to content

Update kron tests to explicitly cover static and dynamic shapes #1912

Open
ayulockedin wants to merge 2 commits intopymc-devs:mainfrom
ayulockedin:add-static-shape-tests
Open

Update kron tests to explicitly cover static and dynamic shapes #1912
ayulockedin wants to merge 2 commits intopymc-devs:mainfrom
ayulockedin:add-static-shape-tests

Conversation

@ayulockedin
Copy link
Contributor

Description

This updates the test suite for pt.linalg.kron to explicitly test both static and non-static shapes. Following reviewer feedback, I added a static_shape parameterization to test_perform and test_kron_commutes_with_inv. This ensures that the operation does not accidentally rely on x.type.shape internally, which would mask bugs and cause silent failures with dynamic shapes in the future. The regression assertions have also been updated to handle tuple lengths dynamically when static shapes are not provided.

Related Issue

Checklist

Type of change

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

@ayulockedin
Copy link
Contributor Author

@ricardoV94 can u take a look at it

[(0, (2, 3), (6, 7)), (1, (2, 3), (4, 3, 5)), (2, (2, 4, 3), (4, 3, 5))],
)
def test_kron_commutes_with_inv(self, i, shp0, shp1):
def test_kron_commutes_with_inv(self, static_shape, i, shp0, shp1):
Copy link
Member

Choose a reason for hiding this comment

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

no reason to touch this test

Comment on lines 768 to 772
# Adjusting the assertion because a non-static tensor won't have integer dimensions in type.shape
if static_shape:
assert kron_xy.type.shape == np_val.shape
else:
assert len(kron_xy.type.shape) == len(np_val.shape)
Copy link
Member

@ricardoV94 ricardoV94 Feb 24, 2026

Choose a reason for hiding this comment

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

define a shape_out like you did shape_x and shape_y, and assert immediately after you define kron_xy that it matches. In the non static shape case it will be (None, None)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants