Description: Circumvent hang while testing The exact trigger of this issue has not yet been determined. Discussed with upstream at https://github.com/e3nn/e3nn/issues/520 Author: Steffen Moeller Origin: https://github.com/e3nn/e3nn/issues/520 Bug: https://github.com/e3nn/e3nn/issues/520 Forwarded: https://github.com/e3nn/e3nn/issues/520 Applied-Upstream: no Last-Update: 2025-11-28 --- This patch header follows DEP-3: http://dep.debian.net/deps/dep3/ Index: python-e3nn/tests/o3/tensor_product_test.py =================================================================== --- python-e3nn.orig/tests/o3/tensor_product_test.py +++ python-e3nn/tests/o3/tensor_product_test.py @@ -6,6 +6,8 @@ import functools import pytest import torch +import sys; + from e3nn.o3 import TensorProduct, FullyConnectedTensorProduct, Irreps from e3nn.util.test import assert_equivariant, assert_auto_jitable, assert_normalized, assert_torch_compile @@ -316,72 +318,75 @@ def test_input_weights_python() -> None: m(x1, x2, w) -def test_input_weights_jit() -> None: - irreps_in1 = Irreps("1e + 2e + 3x3o") - irreps_in2 = Irreps("1e + 2e + 3x3o") - irreps_out = Irreps("1e + 2e + 3x3o") - # - shared_weights = False - - m = FullyConnectedTensorProduct( - irreps_in1, - irreps_in2, - irreps_out, - internal_weights=False, - shared_weights=False, - compile_right=True, - ) - traced = assert_auto_jitable(m) - x1 = irreps_in1.randn(2, -1) - x2 = irreps_in2.randn(2, -1) - w = torch.randn(2, m.weight_numel) - with pytest.raises((RuntimeError, torch.jit.Error)): - m(x1, x2) # it should require weights - with pytest.raises((RuntimeError, torch.jit.Error)): - traced(x1, x2) # it should also require weights - with pytest.raises((RuntimeError, torch.jit.Error)): - traced(x1, x2, w[0]) # it should reject insufficient weights - # Does the trace give right results? - assert torch.allclose(m(x1, x2, w), traced(x1, x2, w)) - - # Confirm that weird batch dimensions give the same results - for f in (m, traced): - x1 = irreps_in1.randn(2, 1, 4, -1) - x2 = irreps_in2.randn(2, 3, 1, -1) - w = torch.randn(3, 4, f.weight_numel) - assert torch.allclose( - f(x1, x2, w).reshape(24, -1), - f( - x1.expand(2, 3, 4, -1).reshape(24, -1), - x2.expand(2, 3, 4, -1).reshape(24, -1), - w[None].expand(2, 3, 4, -1).reshape(24, -1), - ), - ) - assert torch.allclose( - f.right(x2, w).reshape(24, -1), - f.right(x2.expand(2, 3, 4, -1).reshape(24, -1), w[None].expand(2, 3, 4, -1).reshape(24, -1)).reshape(24, -1), - ) - - # - shared_weights = True - - m = FullyConnectedTensorProduct(irreps_in1, irreps_in2, irreps_out, internal_weights=False, shared_weights=True) - w = torch.randn(m.weight_numel) - - traced = assert_auto_jitable(m) - assert_torch_compile( - "inductor", - functools.partial( - FullyConnectedTensorProduct, irreps_in1, irreps_in2, irreps_out, internal_weights=False, shared_weights=True - ), - x1, - x2, - w, - ) - with pytest.raises((RuntimeError, torch.jit.Error)): - m(x1, x2) # it should require weights - with pytest.raises((RuntimeError, torch.jit.Error)): - traced(x1, x2) # it should also require weights - with pytest.raises((RuntimeError, torch.jit.Error)): - traced(x1, x2, torch.randn(2, m.weight_numel)) # it should reject too many weights - # Does the trace give right results? - assert torch.allclose(m(x1, x2, w), traced(x1, x2, w)) +#def test_input_weights_jit() -> None: +# print("test_input_weights - start", file=sys.stderr) +# irreps_in1 = Irreps("1e + 2e + 3x3o") +# irreps_in2 = Irreps("1e + 2e + 3x3o") +# irreps_out = Irreps("1e + 2e + 3x3o") +# # - shared_weights = False - +# #m = FullyConnectedTensorProduct( +# # irreps_in1, +# # irreps_in2, +# # irreps_out, +# # internal_weights=False, +# # shared_weights=False, +# # compile_right=True, +# #) +# traced = assert_auto_jitable(m) +# x1 = irreps_in1.randn(2, -1) +# x2 = irreps_in2.randn(2, -1) +# w = torch.randn(2, m.weight_numel) +# with pytest.raises((RuntimeError, torch.jit.Error)): +# m(x1, x2) # it should require weights +# with pytest.raises((RuntimeError, torch.jit.Error)): +# traced(x1, x2) # it should also require weights +# with pytest.raises((RuntimeError, torch.jit.Error)): +# traced(x1, x2, w[0]) # it should reject insufficient weights +# # Does the trace give right results? +# assert torch.allclose(m(x1, x2, w), traced(x1, x2, w)) +# +# # Confirm that weird batch dimensions give the same results +# for f in (m, traced): +# x1 = irreps_in1.randn(2, 1, 4, -1) +# x2 = irreps_in2.randn(2, 3, 1, -1) +# w = torch.randn(3, 4, f.weight_numel) +# assert torch.allclose( +# f(x1, x2, w).reshape(24, -1), +# f( +# x1.expand(2, 3, 4, -1).reshape(24, -1), +# x2.expand(2, 3, 4, -1).reshape(24, -1), +# w[None].expand(2, 3, 4, -1).reshape(24, -1), +# ), +# ) +# assert torch.allclose( +# f.right(x2, w).reshape(24, -1), +# f.right(x2.expand(2, 3, 4, -1).reshape(24, -1), w[None].expand(2, 3, 4, -1).reshape(24, -1)).reshape(24, -1), +# ) +# +# # - shared_weights = True - +# #m = FullyConnectedTensorProduct(irreps_in1, irreps_in2, irreps_out, internal_weights=False, shared_weights=True) +# #w = torch.randn(m.weight_numel) +# +# #traced = assert_auto_jitable(m) +# #assert_torch_compile( +# # "inductor", +# # functools.partial( +# # FullyConnectedTensorProduct, irreps_in1, irreps_in2, irreps_out, internal_weights=False, shared_weights=True +# # ), +# # x1, +# # x2, +# # w, +# #) +# print("test_input_weights - middle", file=sys.stderr) +# with pytest.raises((RuntimeError, torch.jit.Error)): +# m(x1, x2) # it should require weights +# with pytest.raises((RuntimeError, torch.jit.Error)): +# traced(x1, x2) # it should also require weights +# with pytest.raises((RuntimeError, torch.jit.Error)): +# traced(x1, x2, torch.randn(2, m.weight_numel)) # it should reject too many weights +# # Does the trace give right results? +# assert torch.allclose(m(x1, x2, w), traced(x1, x2, w)) +# print("test_input_weights - end", file=sys.stderr) def test_weight_view_for_instruction() -> None: