[BUG] Attend , dim, MaxVit

[evelynmitchell]

Describe the bug
Tests are failing because Attend doesn't have dim as a parameter.

ERROR zeta/tests/models/test_maxvit.py::test_maxvit_constructor - TypeError: Attend.__init__() got an unexpected keyword argument 'dim'

Attend does not have dim:

class Attend(nn.Module):
    def __init__(
        self,
        *,
        dropout=0.0,
        causal=False,
        heads=None,
        talking_heads=False,
        sparse_topk=None,
        scale=None,
        qk_norm=False,
        flash=False,
        add_zero_kv=False,
        onnxable=False,
    ):

but MaxVit needs it:

class MaxVit(nn.Module):
    def __init__(
        self,
        *,
        num_classes,
        dim,
        depth,
        dim_head: int = 32,
        dim_conv_stem=None,
        window_size: int = 7,
        mbconv_expansion_rate: int = 4,
        mbconv_shrinkage_rate=0.25,
        dropout=0.01,
        channels=3,
    ):

One way to fix this would be to add dim to Attend, but that requires a bit of logic to make it an optional parameter.

Another way to fix it would be to make a multi-model Attend, which would be called by multi-modal models, separate from the sequential attend. (My preferred solution, but may be hard to implement.)

Upvote & Fund

• We're using Polar.sh so you can upvote and help fund this issue.
• We receive the funding once the issue is completed & confirmed by you.
• Thank you in advance for helping prioritize & fund our backlog.

[github-actions]

Stale issue message