code: def test_average():
      params1 = {
          "key1": torch.tensor([1.0, 1.0]),
          "key2": torch.tensor([[2.0, 2.0], [2.0, 2.0]]),
      }
      params2 = {
          "key1": torch.tensor([2.0, 2.0]),
          "key2": torch.tensor([[4.0, 4.0], [4.0, 4.0]]),
      }
      params_average = {
          "key1": torch.tensor([1.5, 1.5]),
          "key2": torch.tensor([[3.0, 3.0], [3.0, 3.0]]),
      }
  
      with tempfile.TemporaryDirectory(
          prefix="test_average"
      ) as work_dir:
          params1_fname = os.path.join(work_dir, "params1")
          params2_fname = os.path.join(work_dir, "params2")
          torch.save(params1, params1_fname)
          torch.save(params2, params2_fname)
  
          # "Average" one params file
          for k, v in average.average([params1_fname]).items():
              assert torch.allclose(v, params1[k])
  
          # Average four params files (two instances each of two files)
          for k, v in average.average(
              [
                  params1_fname,
                  params1_fname,
                  params2_fname,
                  params2_fname,
              ]
          ).items():
              assert torch.allclose(v, params_average[k])
  
file_path: sockeye/test/unit/test_average.py
imports: import os
  import tempfile
  import pytest
  import torch
  import sockeye.average
label_desc: Use image node_id 4 for calling a global function with example usage: test_average() without return types