training# Experimental API. See the NNX page for more details. Metrics Metric Metric.__init__() Metric.compute() Metric.reset() Metric.update() Average Average.__init__() Average.compute() Average.reset() Average.update() Accuracy Accuracy.update() Welford Welford.__init__() Welford.compute() Welford.reset() Welford.update() MultiMetric MultiMetric.__init__() MultiMetric.compute() MultiMetric.reset() MultiMetric.update() Optimizer Optimizer Optimizer.step Optimizer.model Optimizer.tx Optimizer.opt_state Optimizer.__init__() Optimizer.update()