stanford-crfm / haliax

Named Tensors for Legible Deep Learning in JAX
Apache License 2.0
149 stars 11 forks source link

NN Layers: Add GroupNorm #24

Open dlwh opened 1 year ago

dlwh commented 1 year ago

See our LayerNorm and Eqx's groupnorm https://docs.kidger.site/equinox/api/nn/normalisation/#equinox.nn.GroupNorm