linear/conv op dispatch

[ASEM000]

• Enable defining new rules for different weight types other than jax.Array, (e.g. LoRAWeight) on parameter-level of the linear/conv operations, rather than Class-level (e.g. LoRALinear).

Pros

• Parameter level approach enables better composition with at surgery.

  import serket as sk
  import functools as ft
  import re
  net = ...
  
  class CustomWeight(sk.TreeClass):
       ...
  
  @sk.nn.linear.def_type(CustomWeight)
  def _(....):
      # how to handle CustomWeight for linear
      # all layers (and their composition) depending on sk.nn.Linear/sk.nn.linear will be supported by default
  
  # select weights of all linear layers, and apply the CustomWeight
  net = net.at[re.compile("linear_*.")]["weight"].apply(ft.partial(CustomWeight, ...))

• Better Debugging since everything is on the python level

Cons:

Less general than tools operates at jax-level ({initial,final} style) like lorax that can operate on jax progams. However its risky to develop anything with non-public jax API, additionally staged out initial style is harder to debug.

Lets give this a try.