dilated conv layer that can be injected into flax.nn.Conv

google / aqt

Apache License 2.0

247 stars 25 forks source link

dilated conv layer that can be injected into flax.nn.Conv #664

Open Jconn opened 2 months ago

Jconn commented 2 months ago

In aqt_flax, I see AqtDotGeneral but I see no interface for conv general dilated. Is there any technical limitation towards making a AqtConvGeneralDilated for convolutions?