What is replace in TrainState

[pablo2909]

Hello,

I am looking at the code of TrainState in flax.training.train_state.py. The function apply_gradients calls self.replace that is not defined in TrainState, but in PyTreeNode that only raises a NotImplementedError. However when we run it, it does not raise it.
Am I missing something obvious ? I am not sure how this works.

Thank you very much

[marcvanzee]

PyTreeNode implements the init_subclass method, which will converts the subclass (in this case TrainState) into a Flax dataclass, which does implement the replace method.

We inherit from PyTreeNode to avoid type checking issues when using pytype.

[pablo2909]

I see thank you very much !