trax Inconsistency in function's doc-string

Inconsistency in function's doc-string

Open jaindhairyahere opened this issue 1 year ago • 1 comments

Inconsistent function description at https://github.com/google/trax/blob/master/trax/layers/attention.py#L330C1-L342C23

The function states that it "Returns attention-computed per-head activations and unchanged mask." but returns only the activations.

Jan 13 '24 17:01 jaindhairyahere

You are right. Can I raise PR to fix it ?

Apr 18 '24 10:04 punitvara