Makes Metric/Collection inherit from PyNodeTree.

[copybara-service[bot]]

Makes Metric/Collection inherit from PyNodeTree.

That avoids having to decorate Metric/Collection subclasses with @flax.struct.dataclass and removes one class of confusing errors.

Note that all calling code has to be updated accordingly, because subclasses of PyNodeTree can NOT be decorated with @flax.struct.dataclass anymore.