made resizing operations compatible with None

[sayakpaul]

@amyeroberts

Things are progressing in a new good direction, it seems.

I open-sourced the implementation and porting repository. Hence this PR.

I could make the resizing layers handle arbitrary shapes with your advice and guidance. But as you had correctly pointed out, things are still breaking for ResidualSplitHeadMultiAxisGmlpLayer (here).

I guess we might need to rewrite ResidualSplitHeadMultiAxisGmlpLayer subclassing from tf.keras.Model.

WDYT?

Thank you for all your help so far. I really appreciate it.

[sayakpaul]

@amyeroberts, some more findings here.

Refer to this layer once:

maxim-tf/maxim/blocks/misc_gating.py

Line 88 in 6364fd2

def GetSpatialGatingWeights(

It has similar calculations for which (None, None, 3) as an input resolution will break things (here and here). But if I try to make this block a subclass of tf.keras.Model, there are some critical challenges:

• The number of hidden units for this Dense layer is dependent on an intermediate block calculation (here). Same goes for block MLP weights too (here).
• If we initialize these Dense layers inside the call() method of tf.keras.Model, we won't be able to port the pre-trained parameters from the available JAX models appropriately. There are multiple cases like this in this module.

Any further suggestions based on these observations?

Let me know if anything is unclear.