Skip to content
New issue

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

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

Already on GitHub? # to your account

Fix tests #62

Merged
merged 1 commit into from
Jan 13, 2025
Merged

Fix tests #62

merged 1 commit into from
Jan 13, 2025

Conversation

mariogeiger
Copy link
Collaborator

No description provided.

@@ -48,8 +45,8 @@ def make_descriptors():
)

# These ETPs will trigger the symmetricContraction kernel
yield descriptors.spherical_harmonics(cue.SO3(1), [0, 1, 2, 3, 4, 5])
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Here I put back a lower degree polynomial.
The degree 5 was giving a bit too high errors compared to all the other ETP tested.

Comment on lines +186 to +205
@pytest.mark.parametrize("batch_size", [0, 5])
@pytest.mark.parametrize("use_fallback", [True, False])
def test_high_degrees(use_fallback: bool, batch_size: int):
if not use_fallback and not torch.cuda.is_available():
pytest.skip("CUDA is not available")

e = cue.descriptors.spherical_harmonics(cue.SO3(1), [0, 1, 2, 3, 4, 5])
m = cuet.EquivariantTensorProduct(
e,
layout=cue.mul_ir,
device=device,
math_dtype=torch.float32,
use_fallback=use_fallback,
)
inputs = [
torch.randn((batch_size, rep.dim), device=device, dtype=torch.float32)
for rep in e.inputs
]
output = m(*inputs)
assert output.shape == (batch_size, e.output.dim)
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Here I separately test if a high degree polynomial goes through.

@borisfom borisfom merged commit 314092e into main Jan 13, 2025
11 checks passed
@mariogeiger mariogeiger deleted the fix-tests branch January 13, 2025 22:50
# for free to join this conversation on GitHub. Already have an account? # to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants