muP over depth

[AlaaKhaddaj]

Hello,

As the last TP paper points, muP works (empirically) over depth too. Does the current repo support that? Basically, set_base_shapes zips a base model and a wider model, and as such can set the p.infshape for each param p. How can this be adapted to support depth?

Thanks!