From db719e73a5d55f8d1434fccc1e1173b51d24623d Mon Sep 17 00:00:00 2001 From: koki3070 Date: Fri, 19 Jun 2026 11:28:40 +0900 Subject: [PATCH 1/2] =?UTF-8?q?fix:=20issue=20#61=20=E2=80=94=20=E5=AF=BE?= =?UTF-8?q?=E8=A7=92=20self-distance=20=E3=82=92=20sqrt=20=E3=81=8B?= =?UTF-8?q?=E3=82=89=E9=99=A4=E5=A4=96=E3=81=97=20NaN=20=E5=8B=BE=E9=85=8D?= =?UTF-8?q?=E3=82=92=E9=98=B2=E3=81=90?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit off-diagonal のみ sqrt を適用し対角は 0 のまま。ad-hoc eps なし。 Co-authored-by: Cursor --- tda_ml/topology.py | 21 +++++++++- tests/test_topology_gradients.py | 72 ++++++++++++++++++++++++++++++++ 2 files changed, 92 insertions(+), 1 deletion(-) create mode 100644 tests/test_topology_gradients.py diff --git a/tda_ml/topology.py b/tda_ml/topology.py index 99efa7b..e7bbebe 100644 --- a/tda_ml/topology.py +++ b/tda_ml/topology.py @@ -122,5 +122,24 @@ def compute_anisotropic_distance_matrix( f"symmetrize must be 'max' or 'min', got {symmetrize!r}" ) + return _sqrt_off_diagonal_only(dist_sq) + + +def _sqrt_off_diagonal_only(dist_sq: torch.Tensor) -> torch.Tensor: + """ + Mahalanobis distances with zero diagonal without sqrt backward on self-distances. + + Diagonal squared distances are identically zero (self-distance). Differentiating + ``sqrt(x)`` at ``x=0`` yields NaN gradients. VR / persistence uses off-diagonal + entries only, so diagonals stay zero and ``sqrt`` is applied only to off-diagonal + ``dist_sq`` entries (never to the self-distance zeros). + """ dist_sq = torch.clamp(dist_sq, min=0.0) - return torch.sqrt(dist_sq) + dist = torch.zeros_like(dist_sq) + n = dist_sq.shape[-1] + offdiag = ~torch.eye(n, device=dist_sq.device, dtype=torch.bool) + if dist_sq.ndim == 3: + dist[:, offdiag] = torch.sqrt(dist_sq[:, offdiag]) + else: + dist[offdiag] = torch.sqrt(dist_sq[offdiag]) + return dist diff --git a/tests/test_topology_gradients.py b/tests/test_topology_gradients.py new file mode 100644 index 0000000..338f034 --- /dev/null +++ b/tests/test_topology_gradients.py @@ -0,0 +1,72 @@ +"""Gradient health for anisotropic distance matrices (issue #61).""" + +import torch + +from tda_ml.topology import compute_anisotropic_distance_matrix + + +def test_anisotropic_distance_diagonal_is_exactly_zero(): + batch_size, num_points = 1, 5 + points = torch.randn(batch_size, num_points, 2, dtype=torch.float32) + params = torch.tensor( + [ + [ + [1.0, 0.5, 0.0], + [1.0, 0.5, 0.0], + [1.0, 0.5, 0.0], + [1.0, 0.5, 0.0], + [1.0, 0.5, 0.0], + ] + ], + dtype=torch.float32, + ) + probs = torch.full((batch_size, num_points), 0.2, dtype=torch.float32) + + dist = compute_anisotropic_distance_matrix( + points, params, probs=probs, symmetrize="max" + ) + diag = dist.diagonal(dim1=-2, dim2=-1) + assert torch.all(diag == 0.0) + assert torch.all(dist[..., ~torch.eye(num_points, dtype=torch.bool)] > 0) + + +def test_anisotropic_distance_backward_finite_at_zero_diagonal(): + """Self-distance diagonal is zero; backward must not emit inf/nan gradients.""" + batch_size, num_points = 1, 5 + points = torch.randn(batch_size, num_points, 2, dtype=torch.float32) + params = torch.tensor( + [ + [ + [1.0, 0.5, 0.0], + [1.0, 0.5, 0.0], + [1.0, 0.5, 0.0], + [1.0, 0.5, 0.0], + [1.0, 0.5, 0.0], + ] + ], + dtype=torch.float32, + requires_grad=True, + ) + probs = torch.full((batch_size, num_points), 0.2, dtype=torch.float32) + + dist = compute_anisotropic_distance_matrix( + points, params, probs=probs, symmetrize="max" + ) + assert torch.all(dist.diagonal(dim1=-2, dim2=-1) == 0.0) + + dist.sum().backward() + assert params.grad is not None + assert torch.isfinite(params.grad).all() + + +def test_anisotropic_distance_backward_no_sqrt_anomaly_on_diagonal(): + with torch.autograd.detect_anomaly(): + points = torch.randn(1, 8, 2, dtype=torch.float32, requires_grad=True) + params = torch.rand(1, 8, 3, dtype=torch.float32, requires_grad=True) + probs = torch.full((1, 8), 0.3, dtype=torch.float32) + dist = compute_anisotropic_distance_matrix( + points, params, probs=probs, symmetrize="max" + ) + dist.sum().backward() + assert torch.isfinite(points.grad).all() + assert torch.isfinite(params.grad).all() From 43a0b98a6cddaab6e9905450a14c57e344cf674a Mon Sep 17 00:00:00 2001 From: koki3070 Date: Mon, 22 Jun 2026 09:13:40 +0900 Subject: [PATCH 2/2] =?UTF-8?q?fix:=20issue=20#61=20=E2=80=94=20=E3=82=BC?= =?UTF-8?q?=E3=83=AD=E8=B7=9D=E9=9B=A2=E3=82=92=20sqrt=20=E3=81=8B?= =?UTF-8?q?=E3=82=89=E9=99=A4=E5=A4=96=E3=81=97=20coincident=20=E7=82=B9?= =?UTF-8?q?=E3=82=82=E3=82=AB=E3=83=90=E3=83=BC?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit dist_sq > 0 のみ sqrt を適用する構造的マスクに整理し、 厳密一致点・近接点の勾配健全性テストを追加する。 Co-authored-by: Cursor --- tda_ml/topology.py | 13 ++++------ tests/test_topology_gradients.py | 43 ++++++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 8 deletions(-) diff --git a/tda_ml/topology.py b/tda_ml/topology.py index e7bbebe..eb39690 100644 --- a/tda_ml/topology.py +++ b/tda_ml/topology.py @@ -131,15 +131,12 @@ def _sqrt_off_diagonal_only(dist_sq: torch.Tensor) -> torch.Tensor: Diagonal squared distances are identically zero (self-distance). Differentiating ``sqrt(x)`` at ``x=0`` yields NaN gradients. VR / persistence uses off-diagonal - entries only, so diagonals stay zero and ``sqrt`` is applied only to off-diagonal - ``dist_sq`` entries (never to the self-distance zeros). + entries only, so diagonals stay zero and ``sqrt`` is applied only where + ``dist_sq > 0`` (never to zero squared distances, including self-distance and + coincident off-diagonal pairs). """ dist_sq = torch.clamp(dist_sq, min=0.0) dist = torch.zeros_like(dist_sq) - n = dist_sq.shape[-1] - offdiag = ~torch.eye(n, device=dist_sq.device, dtype=torch.bool) - if dist_sq.ndim == 3: - dist[:, offdiag] = torch.sqrt(dist_sq[:, offdiag]) - else: - dist[offdiag] = torch.sqrt(dist_sq[offdiag]) + positive = dist_sq > 0 + dist[positive] = torch.sqrt(dist_sq[positive]) return dist diff --git a/tests/test_topology_gradients.py b/tests/test_topology_gradients.py index 338f034..a06670b 100644 --- a/tests/test_topology_gradients.py +++ b/tests/test_topology_gradients.py @@ -70,3 +70,46 @@ def test_anisotropic_distance_backward_no_sqrt_anomaly_on_diagonal(): dist.sum().backward() assert torch.isfinite(points.grad).all() assert torch.isfinite(params.grad).all() + + +def test_anisotropic_distance_backward_finite_with_coincident_points(): + """Off-diagonal dist_sq is zero when two centers coincide; backward must stay finite.""" + points = torch.tensor( + [[[0.0, 0.0], [1.0, 0.0], [0.0, 0.0]]], + dtype=torch.float32, + requires_grad=True, + ) + params = torch.tensor( + [[[1.0, 0.5, 0.0], [1.0, 0.5, 0.0], [1.0, 0.5, 0.0]]], + dtype=torch.float32, + requires_grad=True, + ) + + dist = compute_anisotropic_distance_matrix(points, params, symmetrize="max") + assert dist[0, 0, 2].item() == 0.0 + assert dist[0, 2, 0].item() == 0.0 + + dist.sum().backward() + assert torch.isfinite(points.grad).all() + assert torch.isfinite(params.grad).all() + + +def test_anisotropic_distance_backward_finite_with_close_points(): + """Near-coincident pairs have small but positive off-diagonal distances.""" + base = torch.randn(1, 6, 2, dtype=torch.float32) + points = base.clone().requires_grad_(True) + with torch.no_grad(): + points[0, 1] = points[0, 0] + 1e-4 + + params = torch.rand(1, 6, 3, dtype=torch.float32, requires_grad=True) + probs = torch.full((1, 6), 0.25, dtype=torch.float32) + + dist = compute_anisotropic_distance_matrix( + points, params, probs=probs, symmetrize="max" + ) + offdiag = ~torch.eye(6, dtype=torch.bool) + assert torch.all(dist[0, offdiag] > 0) + + dist.sum().backward() + assert torch.isfinite(points.grad).all() + assert torch.isfinite(params.grad).all()