Trying to reproduce the results but failing: flax.errors.ScopeParamShapeError

[mashli]

I'm trying to execute the code of this project and I get the following error:

kernel = self.param('kernel', flax.errors.ScopeParamShapeError: Inconsistent shapes between value and initializer for parameter "kernel"

Does anyone know how to solve this problem？

Thank you very much.