Unpack Error when using KFAC with block-diagonal for Dense networks about kfac-jax HOT 4 CLOSED

Ok, so to understand this part you need a bit to read the advanced docs here.

The way KFAC works is in two stages:
It first marks/labels certain computations with specific "layer_tags" and then separately to each tag it assigns a block class.
The generic_tag is only compatible with NaiveDiagonal and NaiveFull curvature blocks, while the Dense blocks are compatible only with the dense_tag currently.

So what needs to happen, is you need to correctly have a dense tag computation in your model in order to be able to use the DenseTwoKroneckerFactored block. If this is not picked up automatically, you can in your code use manually kfac_jax.register_dense to manually register your computation as dense.
I can't comment why it is not picked up automatically, without seeing your model code.

from kfac-jax.

PS: If you resolve your problem please do close the issue tracker.

from kfac-jax.

Ok, so to understand this part you need a bit to read the advanced docs here.

The way KFAC works is in two stages: It first marks/labels certain computations with specific "layer_tags" and then separately to each tag it assigns a block class. The generic_tag is only compatible with NaiveDiagonal and NaiveFull curvature blocks, while the Dense blocks are compatible only with the dense_tag currently.

So what needs to happen, is you need to correctly have a dense tag computation in your model in order to be able to use the DenseTwoKroneckerFactored block. If this is not picked up automatically, you can in your code use manually kfac_jax.register_dense to manually register your computation as dense. I can't comment why it is not picked up automatically, without seeing your model code.

Right, I see. I was using my own defined Dense layers using hk.Module and I assumed that the optimizer would recognize this automatically. My implementation is a bit different from hk.Linear as I was writing my own implementation for KFAC, and variants thereof, and I concatenated the bias vector into the weight matrix. Could this be the cause?

Anyways, I can indeed get it working when I switch out my implementation for the hk.Linear modules. Your comment clears up my confusion! Thanks!

from kfac-jax.

Yes, if as inputs they are separate parameters, but then you concatenate them and use them in that way, KFAC would not automatically recognize it as it does not have the notion that it is a dense layer computation.

from kfac-jax.