diff --git a/tests/ndarray_adapters_test.py b/tests/ndarray_adapters_test.py index 2f38522..537ff0b 100644 --- a/tests/ndarray_adapters_test.py +++ b/tests/ndarray_adapters_test.py @@ -144,15 +144,40 @@ def test_adapter_positional_numpy_consistency(self, array_type, dtype): for fast in (True, False): with self.subTest("summary_fast" if fast else "summary_slow"): summary_info_np = helpers.ensure_text( - np_adapter.get_array_summary(array_np, fast=True) + np_adapter.get_array_summary(array_np, fast=fast) ) summary_info = helpers.ensure_text( - cur_adapter.get_array_summary(array, fast=True) + cur_adapter.get_array_summary(array, fast=fast) ) self.assertEqual( summary_info_np.split(" ", 1)[1], summary_info.split(" ", 1)[1] ) + @parameterized.product( + array_type=["numpy", "jax", "torch"], + values=[ + [1.0, 5.0], + [3.0], + [1.0, 5.0, np.nan], + [1.0, 5.0, np.nan, np.inf, -np.inf], + ], + ) + def test_array_summary_standard_deviation(self, array_type, values): + array_np = np.array(values, dtype=np.float32) + if array_type == "numpy": + array = array_np + elif array_type == "jax": + array = jnp.array(array_np) + else: + array = torch.tensor(array_np) + + adapter = type_registries.lookup_ndarray_adapter(array) + summary = helpers.ensure_text(adapter.get_array_summary(array, fast=False)) + finite_values = array_np[np.isfinite(array_np)] + expected_mean = float(np.mean(finite_values)) + expected_std = float(np.std(finite_values)) + self.assertIn(f" ≈{expected_mean:.2} ±{expected_std:.2}", summary) + def test_pytorch_named_axes_info(self): if not hasattr(torch.Tensor, "rename"): self.skipTest("PyTorch named tensors are not supported in this version") diff --git a/tests/renderer_test.py b/tests/renderer_test.py index e3f7ba9..892d8e9 100644 --- a/tests/renderer_test.py +++ b/tests/renderer_test.py @@ -413,11 +413,11 @@ def hook_that_crashes(node, path, node_renderer): np.arange(3 * 7, dtype=np.float32).reshape((3, 7)) ).requires_grad_(), expected_collapsed=( - """""" + """""" ), expected_expanded=textwrap.dedent( """\ - # torch.Tensor float32(3, 7) ≈1e+01 ±3.7e+01 [≥0.0, ≤2e+01] zero:1 nonzero:20 requires_grad=True + # torch.Tensor float32(3, 7) ≈1e+01 ±6.1 [≥0.0, ≤2e+01] zero:1 nonzero:20 requires_grad=True tensor([[ 0., 1., 2., 3., 4., 5., 6.], [ 7., 8., 9., 10., 11., 12., 13.], [14., 15., 16., 17., 18., 19., 20.]], requires_grad=True)""" @@ -429,11 +429,11 @@ def hook_that_crashes(node, path, node_renderer): torch.tensor(np.arange(3 * 7, dtype=np.float32).reshape((3, 7))) ), expected_collapsed=( - """>""" + """>""" ), expected_expanded=textwrap.dedent( """\ - # torch.Tensor float32(3, 7) ≈1e+01 ±3.7e+01 [≥0.0, ≤2e+01] zero:1 nonzero:20 requires_grad=True grad= + # torch.Tensor float32(3, 7) ≈1e+01 ±6.1 [≥0.0, ≤2e+01] zero:1 nonzero:20 requires_grad=True grad= tensor([[ 0., 1., 2., 3., 4., 5., 6.], [ 7., 8., 9., 10., 11., 12., 13.], [14., 15., 16., 17., 18., 19., 20.]], requires_grad=True)""" @@ -445,11 +445,11 @@ def hook_that_crashes(node, path, node_renderer): torch.tensor(np.arange(3 * 7, dtype=np.float32).reshape((3, 7))) ), expected_collapsed=( - """""" + """""" ), expected_expanded=textwrap.dedent( """\ - # torch.nn.Parameter float32(3, 7) ≈1e+01 ±3.7e+01 [≥0.0, ≤2e+01] zero:1 nonzero:20 + # torch.nn.Parameter float32(3, 7) ≈1e+01 ±6.1 [≥0.0, ≤2e+01] zero:1 nonzero:20 Parameter containing: tensor([[ 0., 1., 2., 3., 4., 5., 6.], [ 7., 8., 9., 10., 11., 12., 13.], @@ -462,10 +462,10 @@ def hook_that_crashes(node, path, node_renderer): torch.tensor(np.arange(3 * 7, dtype=np.float32).reshape((3, 7))) ).requires_grad_(False), expected_collapsed=( - """""" + """""" ), expected_expanded=textwrap.dedent("""\ - # torch.nn.Parameter float32(3, 7) ≈1e+01 ±3.7e+01 [≥0.0, ≤2e+01] zero:1 nonzero:20 requires_grad=False + # torch.nn.Parameter float32(3, 7) ≈1e+01 ±6.1 [≥0.0, ≤2e+01] zero:1 nonzero:20 requires_grad=False Parameter containing: tensor([[ 0., 1., 2., 3., 4., 5., 6.], [ 7., 8., 9., 10., 11., 12., 13.], @@ -602,9 +602,9 @@ def hook_that_crashes(node, path, node_renderer): " foo=, bar=torch.tensor([0., 0., 0., 0., 0.])," " linear=Linear(in_features=10, out_features=10, training=True," - " weight=, bias=, )," + " float32(10,) ≈0.008 ±0.2 [≥-0.31, ≤0.31] nonzero:10>, )," " mod_list=ModuleList(training=True, (0):" " LayerNorm(normalized_shape=(10,), eps=1e-05," " elementwise_affine=True, training=True," @@ -621,7 +621,7 @@ def hook_that_crashes(node, path, node_renderer): # Buffers: bar=torch.tensor([0., 0., 0., 0., 0.]), # Child modules: - linear=Linear(in_features=10, out_features=10, training=True, weight=, bias=, ), + linear=Linear(in_features=10, out_features=10, training=True, weight=, bias=, ), mod_list=ModuleList(training=True, (0): LayerNorm(normalized_shape=(10,), eps=1e-05, elementwise_affine=True, training=True, weight=, bias=, ), (1): SiLU(inplace=False, training=True, ), ), )"""), expected_roundtrip=textwrap.dedent("""\ @@ -632,7 +632,7 @@ def hook_that_crashes(node, path, node_renderer): # Buffers: bar=torch.tensor([0., 0., 0., 0., 0.]), # Child modules: - linear=, bias=, )>, + linear=, bias=, )>, mod_list=, bias=, )>, (1): , )>, )>"""), ), @@ -657,8 +657,8 @@ def hook_that_crashes(node, path, node_renderer): linear=Linear( in_features=10, out_features=10, training=True, # Parameters: - weight=, - bias=, + weight=, + bias=, ), mod_list=ModuleList( training=True, @@ -939,7 +939,7 @@ def hook_that_crashes(node, path, node_renderer): expected_collapsed="[]", expected_expanded=textwrap.dedent("""\ [ - , + , ]"""), ), ) diff --git a/treescope/external/torch_support.py b/treescope/external/torch_support.py index faf05b9..84d0a75 100644 --- a/treescope/external/torch_support.py +++ b/treescope/external/torch_support.py @@ -230,7 +230,7 @@ def get_array_summary( any_finite = torch.any(isfinite) inf_to_nan = torch.where(isfinite, array, torch.nan) mean = torch.nanmean(inf_to_nan) - std = torch.nanmean(torch.square(inf_to_nan - mean)) + std = torch.sqrt(torch.nanmean(torch.square(inf_to_nan - mean))) if any_finite: summary_parts.append(f" ≈{float(mean):.2} ±{float(std):.2}")