Problems with Dense for Linear regression

[RodrigoAVargasHdz]

Hi,

I am trying to optimize linear regression and NN using FLAX. Both models were created with the Dense function from flax.linen.
However, even for a simple 1D linear regression, I haven't managed to converge to an optimal solution.
I created a repo that contains my example, https://github.com/RodrigoAVargasHdz/lr_flax.git

Could you please point me to what are my mistakes?

[marcvanzee]

Hi @RodrigoAVargasHdz, I need a little bit more information to help you here:

• Does it run successfully?
• How did you determine it doesn't converge to an optimal solution? In your code it seems like you aren't printing any losses.
• Did you try tuning the hyperparameters? How did you determine the current ones?