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

[TuanNguyen27]

Thanks for creating such a nice library!

My example use case in flax:

nn = ... # flax.nn.Module
 _, nn_params = nn.init(rng_key, data)
model = flax.nn.Model(nn, nn_params) 

With this, i can call model(x) and i can also access current params via model.params.

For haiku

nn = ... #hk.Module 
nn_params = nn.init(rng_key, data)
# i can do this
partial_fun = lambda data: nn.apply(nn_params, rng_key, data)

I can do partial_fun(x) but because partial_fun is a lambda, i can't access nn_params from partial_fun. Wondering if there's any workaround to achieve this in haiku.

For more context, I am trying to integrate haiku with numpyro so that we can convert traditional NN into Bayesian NN. You can see this issue pyro-ppl/numpyro#705 for more context, or this example notebook

[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 ?

[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.

[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)