Parametrize `attention_normalize_fn` (utils.segment_softmax) by default - Githubissues

google-deepmind / jraph

A Graph Neural Network Library in Jax

https://jraph.readthedocs.io/en/latest/

Apache License 2.0

1.36k stars 88 forks source link

Parametrize `attention_normalize_fn` (utils.segment_softmax) by default #6

Closed sooheon closed 3 years ago

sooheon commented 3 years ago

Allows for using methods other than softmax