code: def amp_metrics_test(trial_class, training_metrics, agg_freq=1):
      loss_prev = None
      GROWTH_INTERVAL = trial_class._growth_interval
      MIN_SCALED_LOSS_TO_REDUCE_SCALE = 32760
      growth_countdown = GROWTH_INTERVAL
      # Only attempt assertions up to and including the penultimate batch, because
      #  we may not have the updated scale from the final batch.
      for idx, (metrics, next_metrics) in enumerate(
          zip(training_metrics[:-1], training_metrics[1:])
      ):
          if (idx + 1) % agg_freq != 0:
              # Only test batches which land on aggregation_frequency boundaries.
              continue
          # Because the scaler is updated during the optimizer step, which occurs after training from
          #  a batch, the metrics dictionary may not have the updated scale, but we can get it
          #  from the next step.
          scale = next_metrics["scale_before"]
          if "scale" in metrics:
              # In cases where we do know the scale immediately after it's updated, we might as well
              #  do this check. If this fails, something is very wrong.
              assert (
                  metrics["scale"] == scale
              ), "scale is inconsistent between batches"
          else:
              metrics["scale"] = scale
          loss = metrics["loss"]
          scale_before = metrics["scale_before"]
          scaled_loss = loss * scale_before
          scale = metrics["scale"]
          growth_countdown -= 1
          if scaled_loss >= MIN_SCALED_LOSS_TO_REDUCE_SCALE:
              assert scale < scale_before, (
                  f"scale was expected to reduce from {scale_before} but did not "
                  f"(scaled_loss={scaled_loss} >= {MIN_SCALED_LOSS_TO_REDUCE_SCALE})) "
              )
              growth_countdown = GROWTH_INTERVAL
          elif growth_countdown == 0:
              assert scale > scale_before, (
                  f"scale was expected to grow but did not "
                  f"(growth_countdown={growth_countdown}) "
              )
              growth_countdown = GROWTH_INTERVAL
          else:
              assert scale == scale_before, (
                  f"scale changed from {scale_before} to {scale} but not when expected "
                  f"(growth_countdown={growth_countdown}) "
                  f"(scaled_loss={scaled_loss} < {MIN_SCALED_LOSS_TO_REDUCE_SCALE})) "
              )
              # Check the accuracy of the gradient change.
              metric_keyname_pairs = [
                  ("loss", "loss_exp"),
                  ("w_after", "w_exp"),
              ]
              if metrics["stage"] in ["small", "zero"]:
                  metric_keyname_pairs.append(("w_before", "w_after"))
              trial_class.check_batch_metrics(
                  metrics,
                  idx,
                  metric_keyname_pairs=metric_keyname_pairs,
                  atol=1e-4,
              )
              if loss_prev is not None and metrics["stage"] == "one":
                  assert (
                      loss <= loss_prev
                  ), "loss was expected to decrease monotonically"
                  loss_prev = loss
  
imports: import contextlib
  import importlib
  import io
  import os
  import pathlib
  import sys
  import typing
  from unittest import mock
  import numpy
  import pytest
  import torch
  from _pytest import monkeypatch
  from torch.distributed import launcher
  import determined
  from determined import pytorch
  from tests.experiment import pytorch_utils, utils
  from tests.experiment.fixtures import pytorch_onevar_model
label_desc: Use image node_id 3 for calling a global function with example usage: amp_metrics_test(trial_class, training_metrics, agg_freq) without return types