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

Compatibility with jit.script and torch.compile: COMPLETE #40

Merged
merged 55 commits into from
Dec 9, 2024
Merged
Changes from 1 commit
Commits
Show all changes
55 commits
Select commit Hold shift + click to select a range
ba9580a
test and quick fix for zero batch
mariogeiger Nov 20, 2024
0bfada9
trigger uniform 1d in test
mariogeiger Nov 20, 2024
fd097c6
satisfy linter
mariogeiger Nov 21, 2024
251fc4d
from typing import
mariogeiger Nov 21, 2024
3498a32
determine math_dtype earlier
mariogeiger Nov 21, 2024
7f3cf05
warning with pip commands
mariogeiger Nov 21, 2024
2624335
remove unused argument
mariogeiger Nov 21, 2024
91f7fce
changelog
mariogeiger Nov 21, 2024
4401048
list of inputs
mariogeiger Nov 21, 2024
ad2db8d
add Fixed subtite
mariogeiger Nov 21, 2024
dca96a8
Merge branch 'zero-batch' into list-inputs
mariogeiger Nov 21, 2024
889051a
changelog
mariogeiger Nov 21, 2024
c23816a
Merge branch 'main' into list-inputs
mariogeiger Nov 21, 2024
0487d77
Merge branch 'main' into list-inputs
mariogeiger Dec 3, 2024
bc6b405
add test for torch.jit.script
mariogeiger Dec 3, 2024
c8de185
fix
mariogeiger Dec 3, 2024
5e00b37
Merge branch 'list-inputs' into jit-script
mariogeiger Dec 3, 2024
16e4450
remove keyword-only and import in the forward
mariogeiger Dec 3, 2024
e979b0f
Merge branch 'main' into jit-script
mariogeiger Dec 4, 2024
b2c4fbb
low lvl script tests
mariogeiger Dec 4, 2024
4669a86
TensorProduct working with script()
borisfom Dec 4, 2024
dc9d5b0
add 4 operands tests
mariogeiger Dec 4, 2024
334b460
Unit tests run
borisfom Dec 5, 2024
79e7c5f
Restoring debug logging
borisfom Dec 5, 2024
46a0478
Merge branch 'jit-script' into jit-script
borisfom Dec 5, 2024
401fd53
Merge branch 'jit-script' of github.com:NVIDIA/cuEquivariance into ji…
borisfom Dec 5, 2024
8fce54b
Merge remote-tracking branch 'b/jit-script' into jit-script
borisfom Dec 5, 2024
6c5cdb0
Parameterized script test
borisfom Dec 5, 2024
e21c45f
Fixed transpose for script(), script_test successful
borisfom Dec 5, 2024
779dd9c
Fixed input mutation
borisfom Dec 5, 2024
c315857
Fixed tests
borisfom Dec 6, 2024
ab590c8
format with black
mariogeiger Dec 6, 2024
ec1eb27
format with black
mariogeiger Dec 6, 2024
faf235e
fix tests
mariogeiger Dec 6, 2024
c476af9
fix missing parenthesis
mariogeiger Dec 6, 2024
994b8d9
fix tests: increase torch._dynamo.config.cache_size_limit
mariogeiger Dec 6, 2024
f240eb8
fix docstring tests
mariogeiger Dec 6, 2024
fbfb9d0
replace == by is
mariogeiger Dec 6, 2024
dc20be5
clean use_fallback conditions
mariogeiger Dec 6, 2024
4b201c3
fix
mariogeiger Dec 6, 2024
b5b59b8
fix
mariogeiger Dec 6, 2024
72baf17
Export test added, scripting fallback attempt
borisfom Dec 7, 2024
5a94b09
Merge remote-tracking branch 'b/jit-script' into jit-script
borisfom Dec 7, 2024
6bdf924
Merge branch 'main' into jit-script
mariogeiger Dec 9, 2024
8d31929
enable tests on cpu
mariogeiger Dec 9, 2024
8afa056
fix tests
mariogeiger Dec 9, 2024
09bbc8d
fix ruff
mariogeiger Dec 9, 2024
9c38168
fix
mariogeiger Dec 9, 2024
de9af8f
fix docstring tests
mariogeiger Dec 9, 2024
999a31d
add -x to tests
mariogeiger Dec 9, 2024
905e716
changelog
mariogeiger Dec 9, 2024
975e9c8
test
mariogeiger Dec 9, 2024
093e8e4
move utils into test file
mariogeiger Dec 9, 2024
2712f54
fix
mariogeiger Dec 9, 2024
008ee3d
Merge branch 'main' into jit-script
mariogeiger Dec 9, 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
Prev Previous commit
Next Next commit
fix
  • Loading branch information
mariogeiger committed Dec 9, 2024
commit 9c38168cbc5f854ecaa579319b8825e395fdb58a
Original file line number Diff line number Diff line change
@@ -70,7 +70,7 @@ class FullyConnectedTensorProductConv(nn.Module):
having 16 channels. edge_emb.size(1) must match the size of the input layer: 6

>>> conv1 = FullyConnectedTensorProductConv(in_irreps, sh_irreps, out_irreps,
... mlp_channels=[6, 16, 16], mlp_activation=nn.ReLU(), layout=cue.ir_mul).cuda()
... mlp_channels=[6, 16, 16], mlp_activation=nn.ReLU(), layout=cue.ir_mul)
>>> conv1
FullyConnectedTensorProductConv(...)
>>> # out = conv1(src_features, edge_sh, edge_emb, graph)
@@ -92,7 +92,7 @@ class FullyConnectedTensorProductConv(nn.Module):
**Case 3**: No MLP, edge_emb will be directly used as the tensor product weights:

>>> conv3 = FullyConnectedTensorProductConv(in_irreps, sh_irreps, out_irreps,
... mlp_channels=None, layout=cue.ir_mul).cuda()
... mlp_channels=None, layout=cue.ir_mul)
>>> # out = conv3(src_features, edge_sh, edge_emb, graph)
"""