refactor dot_product_attention to be more general but also check argu… - Githubissues

stanford-crfm / haliax

Named Tensors for Legible Deep Learning in JAX

Apache License 2.0

140 stars 9 forks source link

refactor dot_product_attention to be more general but also check argu… #49

Closed dlwh closed 8 months ago

dlwh commented 8 months ago

…ments better. add self_attention function