Make ASAPooling jittable

[rusty1s]

No description provided.

[codecov]

Codecov Report

Merging #5395 (54310df) into master (7cea70d) will decrease coverage by 1.99%.
The diff coverage is 78.57%.

@@            Coverage Diff             @@
##           master    #5395      +/-   ##
==========================================
- Coverage   85.33%   83.34%   -2.00%     
==========================================
  Files         342      342              
  Lines       18774    18782       +8     
==========================================
- Hits        16020    15653     -367     
- Misses       2754     3129     +375     

Impacted Files	Coverage Δ
torch_geometric/nn/pool/asap.py	92.10% <71.42%> (-7.90%)	⬇️
torch_geometric/nn/pool/topk_pool.py	100.00% <100.00%> (ø)
torch_geometric/nn/models/dimenet_utils.py	0.00% <0.00%> (-75.52%)	⬇️
torch_geometric/nn/models/dimenet.py	14.51% <0.00%> (-53.00%)	⬇️
torch_geometric/profile/profile.py	37.89% <0.00%> (-26.32%)	⬇️
torch_geometric/nn/conv/utils/typing.py	81.25% <0.00%> (-17.50%)	⬇️
torch_geometric/nn/inits.py	67.85% <0.00%> (-7.15%)	⬇️
torch_geometric/transforms/add_self_loops.py	94.44% <0.00%> (-5.56%)	⬇️
torch_geometric/nn/resolver.py	88.88% <0.00%> (-5.56%)	⬇️
torch_geometric/sampler/neighbor_sampler.py	93.28% <0.00%> (-5.37%)	⬇️
... and 11 more

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more