Repair Liskov error with attention modules

[NeilGirdhar]

trafficstars

Fixes # (issue)

Checklist

• [x] This PR fixes a minor issue (e.g.: typo or small bug) or improves the docs (you can dismiss the other checks if that's the case).

[NeilGirdhar]

This is a Liskov error because method overrides should accept a superset of whatever their overridden methods accept. Otherwise, the child class cannot be used polymorphically.

[codecov-commenter]

Codecov Report

Merging #1962 (1a48c12) into main (34790f5) will increase coverage by 0.32%. The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main    #1962      +/-   ##
==========================================
+ Coverage   70.97%   71.30%   +0.32%     
==========================================
  Files          58       59       +1     
  Lines        4931     4987      +56     
==========================================
+ Hits         3500     3556      +56     
  Misses       1431     1431              

Impacted Files	Coverage Δ
flax/linen/attention.py	95.86% <100.00%> (+0.14%)	:arrow_up:
flax/linen/__init__.py	100.00% <0.00%> (ø)
flax/linen/combinators.py	100.00% <0.00%> (ø)
flax/linen/linear.py	99.14% <0.00%> (+<0.01%)	:arrow_up:
flax/linen/transforms.py	96.05% <0.00%> (+0.16%)	:arrow_up:
flax/core/lift.py	95.74% <0.00%> (+0.26%)	:arrow_up:
flax/struct.py	74.62% <0.00%> (+1.61%)	:arrow_up:

---

Continue to review full report at Codecov.

Legend - Click here to learn more Δ = absolute <relative> (impact), ø = not affected, ? = missing data Powered by Codecov. Last update 34790f5...1a48c12. Read the comment docs.

[NeilGirdhar]

This seems like quite a lot of duplication, and feels a bit odd since the actual attributes aren't defined in any of those classes

I agree. Ideally RTD should collate the attributes of all parent classes.

We will also have to add docstring to both call functions since these are no longer annotated either -

Right, but shouldn't that be done anyway? They have slightly different parameters.

I am not sure whether this will make things better or worse...

The problem with not fixing it is that someone may expect a container of these objects to work:

l: list[MultiHeadDotProductAttention]
l.append(MultiHeadDotProductAttention(...))
l.append(SelfAttention(...))
for x in l:
  x(...)  # passes type check, but fails

Inheritance is fine to share behavior between these two classes, but there isn't an "is-a" relationship between them. At least theoretically, this is a misuse of inheritance.