In haiku, is there something equivalent to flax.nn.Model?

[google-deepmind]

In haiku, is there something equivalent to flax.nn.Model?

[cgarciae]

@TuanNguyen27 From the base implementation of the Module class + hooks (get_parameter and friends) probably not. If you need any layer not yet ported its usually very straight forward since the code is identical 99% of the time; you can post an issue since its our goal to have all of them.

[cgarciae]

Haiku doesn't offers a Model interface. You can try Elegy which offers a Keras-like Model interface, however since version 0.2.0 it is no longer compatible with Haiku, the Module system is very similar though.

[TuanNguyen27]

thanks for the tip, @cgarciae ! Just for my own curiosity, is there any feature available in haiku but not in Elegy ?

[tomhennigan]

We don't ship a model interface in Haiku (in general we try to be un-opinionated wrt how you use your neural network), but you can write a small wrapper class if you want to hold your apply function and parameters as an object:

@dataclass
class Model:
  params: hk.Params
  apply: Callable

  def __call__(self, x):
    return self.apply(self.params, x)

model = Model(params, f.apply)
print(model.params)
out = model(x)