brainets / hoi

Higher-Order Interactions

https://brainets.github.io/hoi/

BSD 3-Clause "New" or "Revised" License

18 stars 7 forks source link

Closed EtienneCmb closed 1 year ago

EtienneCmb commented 1 year ago

Goal

Comparison between the loop-free tensor implementation of the entropy vs. using jax.vmap looping over the first axis.

EtienneCmb commented 1 year ago

Vmap scales better with larger data

EtienneCmb commented 1 year ago

Same on cpu