code: def test_stats_of_dict():
      stats = StatsOfDict(
          types={"aaa": CountMeanStats}, default=CountStats
      )
      stats.update({"aaa": torch.tensor(0.0)})
      stats.update({"aaa": torch.tensor(1.0), "bbb": torch.randn(3, 3)})
      stats.update({"aaa": torch.tensor(2.0), "bbb": torch.randn(3, 3)})
      actual = stats.get()
  
      expected = {
          "aaa": {"count": 3, "mean": torch.tensor(1.0)},
          "bbb": {"count": 2},
      }
      assert_close(actual, expected)
  
file_path: pyro/tests/ops/test_streaming.py
imports: import functools
  import pytest
  import torch
  from pyro.ops.streaming import CountMeanStats, CountMeanVarianceStats, CountStats, StackStats, StatsOfDict
  from tests.common import assert_close
label_desc: Use image node_id 5 for calling a global function with example usage: test_stats_of_dict() without return types