search

google / trax

Trax — Deep Learning with Clear Code and Speed
Apache License 2.0
8.01k stars 813 forks source link

Inconsistency in function's doc-string #1794

Open jaindhairyahere opened 5 months ago

jaindhairyahere commented 5 months ago

Description

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.

punitvara commented 2 months ago

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