Visualize comparison detailed results in dump comparator (#19565)
This commit is contained in:
@@ -639,119 +639,5 @@ class TestThdCpConcat:
|
||||
)
|
||||
|
||||
|
||||
class TestThdCpConcat:
|
||||
def test_single_seq(self) -> None:
|
||||
"""Single seq THD unshard: 2 ranks → per-seq concat."""
|
||||
rank0 = torch.tensor([1, 2, 3]).refine_names("t")
|
||||
rank1 = torch.tensor([4, 5, 6]).refine_names("t")
|
||||
|
||||
plan = UnsharderPlan(
|
||||
axis=ParallelAxis.CP,
|
||||
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3]),
|
||||
groups=[[0, 1]],
|
||||
)
|
||||
with warning_sink.context():
|
||||
result = execute_unsharder_plan(plan, [rank0, rank1])
|
||||
|
||||
assert len(result) == 1
|
||||
expected = torch.tensor([1, 2, 3, 4, 5, 6])
|
||||
assert torch.equal(result[0].rename(None), expected)
|
||||
|
||||
def test_multi_seq(self) -> None:
|
||||
"""Multi-seq THD unshard: 2 ranks, seq_lens=[50, 32, 46]."""
|
||||
# rank0: [seqA_r0(50) | seqB_r0(32) | pad_r0(46)]
|
||||
# rank1: [seqA_r1(50) | seqB_r1(32) | pad_r1(46)]
|
||||
seq_a_r0 = torch.arange(0, 50)
|
||||
seq_b_r0 = torch.arange(100, 132)
|
||||
pad_r0 = torch.full((46,), -1)
|
||||
rank0 = torch.cat([seq_a_r0, seq_b_r0, pad_r0]).refine_names("t")
|
||||
|
||||
seq_a_r1 = torch.arange(50, 100)
|
||||
seq_b_r1 = torch.arange(132, 164)
|
||||
pad_r1 = torch.full((46,), -2)
|
||||
rank1 = torch.cat([seq_a_r1, seq_b_r1, pad_r1]).refine_names("t")
|
||||
|
||||
plan = UnsharderPlan(
|
||||
axis=ParallelAxis.CP,
|
||||
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[50, 32, 46]),
|
||||
groups=[[0, 1]],
|
||||
)
|
||||
with warning_sink.context():
|
||||
result = execute_unsharder_plan(plan, [rank0, rank1])
|
||||
|
||||
assert len(result) == 1
|
||||
unsharded: torch.Tensor = result[0].rename(None)
|
||||
|
||||
# seqA: r0(50) + r1(50) = 100 tokens, values 0..99
|
||||
assert torch.equal(unsharded[:100], torch.cat([seq_a_r0, seq_a_r1]))
|
||||
# seqB: r0(32) + r1(32) = 64 tokens
|
||||
assert torch.equal(unsharded[100:164], torch.cat([seq_b_r0, seq_b_r1]))
|
||||
# pad: r0(46) + r1(46) = 92 tokens
|
||||
assert torch.equal(unsharded[164:256], torch.cat([pad_r0, pad_r1]))
|
||||
|
||||
def test_with_hidden_dim(self) -> None:
|
||||
"""THD unshard with trailing hidden dim: shape [T, H]."""
|
||||
torch.manual_seed(42)
|
||||
hidden: int = 4
|
||||
# rank0: [seqA_r0(3, 4) | seqB_r0(2, 4)]
|
||||
# rank1: [seqA_r1(3, 4) | seqB_r1(2, 4)]
|
||||
seq_a_r0 = torch.randn(3, hidden)
|
||||
seq_b_r0 = torch.randn(2, hidden)
|
||||
rank0 = torch.cat([seq_a_r0, seq_b_r0]).refine_names("t", "h")
|
||||
|
||||
seq_a_r1 = torch.randn(3, hidden)
|
||||
seq_b_r1 = torch.randn(2, hidden)
|
||||
rank1 = torch.cat([seq_a_r1, seq_b_r1]).refine_names("t", "h")
|
||||
|
||||
plan = UnsharderPlan(
|
||||
axis=ParallelAxis.CP,
|
||||
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]),
|
||||
groups=[[0, 1]],
|
||||
)
|
||||
with warning_sink.context():
|
||||
result = execute_unsharder_plan(plan, [rank0, rank1])
|
||||
|
||||
assert len(result) == 1
|
||||
unsharded: torch.Tensor = result[0].rename(None)
|
||||
|
||||
assert unsharded.shape == (10, hidden)
|
||||
assert torch.equal(unsharded[:6], torch.cat([seq_a_r0, seq_a_r1]))
|
||||
assert torch.equal(unsharded[6:10], torch.cat([seq_b_r0, seq_b_r1]))
|
||||
|
||||
def test_with_leading_batch_dim(self) -> None:
|
||||
"""THD unshard with leading batch dim: shape [B, T, H], t is dim=1."""
|
||||
torch.manual_seed(42)
|
||||
batch: int = 2
|
||||
hidden: int = 4
|
||||
# rank0: [seqA_r0(3) | seqB_r0(2)] per batch item
|
||||
# rank1: [seqA_r1(3) | seqB_r1(2)] per batch item
|
||||
seq_a_r0 = torch.randn(batch, 3, hidden)
|
||||
seq_b_r0 = torch.randn(batch, 2, hidden)
|
||||
rank0 = torch.cat([seq_a_r0, seq_b_r0], dim=1).refine_names("b", "t", "h")
|
||||
|
||||
seq_a_r1 = torch.randn(batch, 3, hidden)
|
||||
seq_b_r1 = torch.randn(batch, 2, hidden)
|
||||
rank1 = torch.cat([seq_a_r1, seq_b_r1], dim=1).refine_names("b", "t", "h")
|
||||
|
||||
plan = UnsharderPlan(
|
||||
axis=ParallelAxis.CP,
|
||||
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]),
|
||||
groups=[[0, 1]],
|
||||
)
|
||||
with warning_sink.context():
|
||||
result = execute_unsharder_plan(plan, [rank0, rank1])
|
||||
|
||||
assert len(result) == 1
|
||||
unsharded: torch.Tensor = result[0].rename(None)
|
||||
|
||||
assert unsharded.shape == (batch, 10, hidden)
|
||||
# seqA: r0(3) + r1(3) = 6 tokens per batch
|
||||
assert torch.equal(unsharded[:, :6, :], torch.cat([seq_a_r0, seq_a_r1], dim=1))
|
||||
# seqB: r0(2) + r1(2) = 4 tokens per batch
|
||||
assert torch.equal(
|
||||
unsharded[:, 6:10, :], torch.cat([seq_b_r0, seq_b_r1], dim=1)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
|
||||
@@ -1019,80 +1019,6 @@ class TestEntrypointAxisSwapper:
|
||||
assert comp.name == "hidden"
|
||||
|
||||
|
||||
class TestEntrypointAxisSwapper:
|
||||
"""Test cross-framework dim reordering through the full entrypoint pipeline."""
|
||||
|
||||
def test_axis_swap_different_dim_order(self, tmp_path, capsys):
|
||||
"""Baseline dims 'b h d' vs target dims 'b d h': axis swapper rearranges baseline to match."""
|
||||
torch.manual_seed(42)
|
||||
full_tensor = torch.randn(4, 8, 16)
|
||||
|
||||
baseline_dir = tmp_path / "baseline"
|
||||
target_dir = tmp_path / "target"
|
||||
|
||||
_create_rank_dump(
|
||||
baseline_dir,
|
||||
rank=0,
|
||||
name="hidden",
|
||||
tensor=full_tensor,
|
||||
dims="b h d",
|
||||
)
|
||||
_create_rank_dump(
|
||||
target_dir,
|
||||
rank=0,
|
||||
name="hidden",
|
||||
tensor=full_tensor.permute(0, 2, 1).contiguous(),
|
||||
dims="b d h",
|
||||
)
|
||||
|
||||
args = _make_args(
|
||||
baseline_dir / _FIXED_EXP_NAME,
|
||||
target_dir / _FIXED_EXP_NAME,
|
||||
diff_threshold=1e-3,
|
||||
)
|
||||
|
||||
records = _run_and_parse(args, capsys)
|
||||
comp = _assert_single_comparison_passed(records)
|
||||
assert comp.name == "hidden"
|
||||
assert comp.baseline.shape == [4, 16, 8]
|
||||
assert comp.target.shape == [4, 16, 8]
|
||||
|
||||
def test_axis_swap_with_tp_unshard(self, tmp_path, capsys):
|
||||
"""Baseline TP=2 with dims 'b h(tp) d' vs target TP=2 with dims 'b d h(tp)': unshard + axis swap."""
|
||||
torch.manual_seed(42)
|
||||
full_tensor = torch.randn(4, 8, 16)
|
||||
|
||||
baseline_dir = tmp_path / "baseline"
|
||||
target_dir = tmp_path / "target"
|
||||
|
||||
_create_tp_sharded_dumps(
|
||||
baseline_dir,
|
||||
full_tensor=full_tensor,
|
||||
name="hidden",
|
||||
tp_size=2,
|
||||
shard_dim=1,
|
||||
dims_str="b h(tp) d",
|
||||
)
|
||||
_create_tp_sharded_dumps(
|
||||
target_dir,
|
||||
full_tensor=full_tensor.permute(0, 2, 1).contiguous(),
|
||||
name="hidden",
|
||||
tp_size=2,
|
||||
shard_dim=2,
|
||||
dims_str="b d h(tp)",
|
||||
)
|
||||
|
||||
args = _make_args(
|
||||
baseline_dir / _FIXED_EXP_NAME,
|
||||
target_dir / _FIXED_EXP_NAME,
|
||||
diff_threshold=1e-3,
|
||||
)
|
||||
|
||||
records = _run_and_parse(args, capsys)
|
||||
comp = _assert_single_comparison_passed(records)
|
||||
assert comp.name == "hidden"
|
||||
|
||||
|
||||
class TestEntrypointReplicatedAxis:
|
||||
"""Test replicated-axis scenarios through the full entrypoint pipeline."""
|
||||
|
||||
@@ -1578,6 +1504,52 @@ class TestEntrypointNonTensorValues:
|
||||
assert roundtripped.values_equal is True
|
||||
|
||||
|
||||
# ───────────────────── Visualization integration tests ─────────────────────
|
||||
|
||||
|
||||
class TestEntrypointVisualize:
|
||||
"""Test --visualize-bundle-details integration."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _skip_if_no_matplotlib(self) -> None:
|
||||
pytest.importorskip("matplotlib")
|
||||
|
||||
def test_visualize_creates_pngs(self, tmp_path, capsys):
|
||||
"""--visualize-bundle-details with --filter produces PNG files."""
|
||||
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
|
||||
viz_dir = tmp_path / "viz_out"
|
||||
args = _make_args(
|
||||
baseline_path,
|
||||
target_path,
|
||||
grouping="raw",
|
||||
filter="tensor_a",
|
||||
viz_bundle_details=True,
|
||||
viz_output_dir=str(viz_dir),
|
||||
)
|
||||
|
||||
records = _run_and_parse(args, capsys)
|
||||
assert len(_get_comparisons(records)) == 1
|
||||
|
||||
png_files = list(viz_dir.glob("*.png"))
|
||||
assert len(png_files) == 1
|
||||
assert png_files[0].stat().st_size > 0
|
||||
|
||||
def test_no_visualize_no_png(self, tmp_path, capsys):
|
||||
"""Without --visualize-bundle-details, no PNGs are created."""
|
||||
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"])
|
||||
viz_dir = tmp_path / "viz_out"
|
||||
args = _make_args(
|
||||
baseline_path,
|
||||
target_path,
|
||||
grouping="raw",
|
||||
viz_bundle_details=False,
|
||||
viz_output_dir=str(viz_dir),
|
||||
)
|
||||
|
||||
_run_and_parse(args, capsys)
|
||||
assert not viz_dir.exists() or len(list(viz_dir.glob("*.png"))) == 0
|
||||
|
||||
|
||||
# --------------------------- Assertion helpers -------------------
|
||||
|
||||
|
||||
@@ -1702,6 +1674,8 @@ def _make_args(baseline_path: Path, target_path: Path, **overrides) -> Namespace
|
||||
filter=None,
|
||||
output_format="json",
|
||||
grouping="logical",
|
||||
viz_bundle_details=False,
|
||||
viz_output_dir="/tmp/comparator_viz/",
|
||||
)
|
||||
defaults.update(overrides)
|
||||
return Namespace(**defaults)
|
||||
|
||||
203
test/registered/debug_utils/comparator/test_manually_verify.py
Normal file
203
test/registered/debug_utils/comparator/test_manually_verify.py
Normal file
@@ -0,0 +1,203 @@
|
||||
"""Visual comparison figure tests — CI sanity check + human verification.
|
||||
|
||||
This file serves two purposes:
|
||||
1. CI sanity check: ensures generate_comparison_figure() runs without errors
|
||||
across various tensor scenarios (registered via register_cpu_ci).
|
||||
2. Human verification: all generated PNGs are copied to /tmp/comparator_manual_verify/
|
||||
so they can be pulled back to a local machine for visual inspection.
|
||||
|
||||
Run:
|
||||
python -m pytest test/registered/debug_utils/comparator/test_manually_verify.py -x -v
|
||||
|
||||
Human verification:
|
||||
After running, images are at /tmp/comparator_manual_verify/.
|
||||
Each test's docstring describes the expected visual appearance.
|
||||
"""
|
||||
|
||||
import shutil
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=60, suite="default", nightly=True)
|
||||
|
||||
_PUBLISH_DIR: Path = Path("/tmp/comparator_manual_verify")
|
||||
_PNG_MAGIC: bytes = b"\x89PNG"
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def publish_dir() -> Path:
|
||||
"""Fixed output dir for human inspection — files are copied here after generation."""
|
||||
if _PUBLISH_DIR.exists():
|
||||
shutil.rmtree(_PUBLISH_DIR)
|
||||
_PUBLISH_DIR.mkdir(parents=True)
|
||||
return _PUBLISH_DIR
|
||||
|
||||
|
||||
def _assert_valid_png(path: Path) -> None:
|
||||
assert path.exists(), f"PNG not created: {path}"
|
||||
assert path.stat().st_size > 0, f"PNG is empty: {path}"
|
||||
with open(path, "rb") as f:
|
||||
magic: bytes = f.read(4)
|
||||
assert magic == _PNG_MAGIC, f"Not a valid PNG: {path}"
|
||||
|
||||
|
||||
def _generate_and_publish(
|
||||
*,
|
||||
baseline: torch.Tensor,
|
||||
target: torch.Tensor,
|
||||
name: str,
|
||||
tmp_path: Path,
|
||||
publish_dir: Path,
|
||||
) -> Path:
|
||||
from sglang.srt.debug_utils.comparator.visualizer import (
|
||||
generate_comparison_figure,
|
||||
)
|
||||
|
||||
output_path: Path = tmp_path / f"{name}.png"
|
||||
generate_comparison_figure(
|
||||
baseline=baseline,
|
||||
target=target,
|
||||
name=name,
|
||||
output_path=output_path,
|
||||
)
|
||||
|
||||
_assert_valid_png(output_path)
|
||||
shutil.copy2(src=output_path, dst=publish_dir / output_path.name)
|
||||
return output_path
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _skip_if_no_matplotlib() -> None:
|
||||
pytest.importorskip("matplotlib")
|
||||
|
||||
|
||||
class TestManuallyVerify:
|
||||
def test_normal_small_diff(self, tmp_path: Path, publish_dir: Path) -> None:
|
||||
"""Two nearly-identical tensors (randn + 0.01 noise).
|
||||
|
||||
Expected: All 6 panel rows visible. Diff heatmap nearly uniform light color.
|
||||
Hist2d tightly clustered along the red diagonal line.
|
||||
"""
|
||||
baseline: torch.Tensor = torch.randn(32, 64)
|
||||
target: torch.Tensor = baseline + torch.randn(32, 64) * 0.01
|
||||
|
||||
_generate_and_publish(
|
||||
baseline=baseline,
|
||||
target=target,
|
||||
name="normal_small_diff",
|
||||
tmp_path=tmp_path,
|
||||
publish_dir=publish_dir,
|
||||
)
|
||||
|
||||
def test_significant_diff(self, tmp_path: Path, publish_dir: Path) -> None:
|
||||
"""Two tensors with larger differences (randn + 0.5 noise).
|
||||
|
||||
Expected: All 6 panel rows visible. Diff heatmap shows noticeable structure.
|
||||
Hist2d scatter is broader, spread away from the diagonal.
|
||||
"""
|
||||
baseline: torch.Tensor = torch.randn(32, 64)
|
||||
target: torch.Tensor = baseline + torch.randn(32, 64) * 0.5
|
||||
|
||||
_generate_and_publish(
|
||||
baseline=baseline,
|
||||
target=target,
|
||||
name="significant_diff",
|
||||
tmp_path=tmp_path,
|
||||
publish_dir=publish_dir,
|
||||
)
|
||||
|
||||
def test_shape_mismatch(self, tmp_path: Path, publish_dir: Path) -> None:
|
||||
"""Baseline 32x64, target 16x32 — shapes do not match.
|
||||
|
||||
Expected: Only 2 panel rows (baseline heatmap, target heatmap).
|
||||
No diff/histogram/hist2d/sampled panels since diff cannot be computed.
|
||||
"""
|
||||
baseline: torch.Tensor = torch.randn(32, 64)
|
||||
target: torch.Tensor = torch.randn(16, 32)
|
||||
|
||||
_generate_and_publish(
|
||||
baseline=baseline,
|
||||
target=target,
|
||||
name="shape_mismatch",
|
||||
tmp_path=tmp_path,
|
||||
publish_dir=publish_dir,
|
||||
)
|
||||
|
||||
def test_large_tensor(self, tmp_path: Path, publish_dir: Path) -> None:
|
||||
"""4000x4000 tensor — triggers internal downsampling.
|
||||
|
||||
Expected: Figure renders normally without OOM. Downsampled panels
|
||||
should still look reasonable.
|
||||
"""
|
||||
baseline: torch.Tensor = torch.randn(4000, 4000)
|
||||
target: torch.Tensor = baseline + torch.randn(4000, 4000) * 0.001
|
||||
|
||||
_generate_and_publish(
|
||||
baseline=baseline,
|
||||
target=target,
|
||||
name="large_tensor",
|
||||
tmp_path=tmp_path,
|
||||
publish_dir=publish_dir,
|
||||
)
|
||||
|
||||
def test_1d_tensor(self, tmp_path: Path, publish_dir: Path) -> None:
|
||||
"""1D tensor (256,) — internally reshaped to 2D before plotting.
|
||||
|
||||
Expected: All 6 panel rows visible. The heatmap shape reflects the
|
||||
reshaped 2D form, not the original 1D.
|
||||
"""
|
||||
baseline: torch.Tensor = torch.randn(256)
|
||||
target: torch.Tensor = baseline + 0.01
|
||||
|
||||
_generate_and_publish(
|
||||
baseline=baseline,
|
||||
target=target,
|
||||
name="1d_tensor",
|
||||
tmp_path=tmp_path,
|
||||
publish_dir=publish_dir,
|
||||
)
|
||||
|
||||
def test_constant_tensor(self, tmp_path: Path, publish_dir: Path) -> None:
|
||||
"""All-zero baseline, tiny-valued target.
|
||||
|
||||
Expected: Colorbar range is extremely small. Histogram concentrates in
|
||||
a single bin. No rendering errors from near-zero variance.
|
||||
"""
|
||||
baseline: torch.Tensor = torch.zeros(32, 64)
|
||||
target: torch.Tensor = torch.ones(32, 64) * 1e-8
|
||||
|
||||
_generate_and_publish(
|
||||
baseline=baseline,
|
||||
target=target,
|
||||
name="constant_tensor",
|
||||
tmp_path=tmp_path,
|
||||
publish_dir=publish_dir,
|
||||
)
|
||||
|
||||
def test_extreme_values(self, tmp_path: Path, publish_dir: Path) -> None:
|
||||
"""Tensor containing values spanning 1e-10 to 1e10.
|
||||
|
||||
Expected: Log10 panels handle the wide range gracefully. No inf/nan
|
||||
artifacts in the rendered figure.
|
||||
"""
|
||||
baseline: torch.Tensor = torch.randn(32, 64).abs()
|
||||
baseline[0, 0] = 1e-10
|
||||
baseline[0, 1] = 1e10
|
||||
target: torch.Tensor = baseline + torch.randn(32, 64) * 0.01
|
||||
|
||||
_generate_and_publish(
|
||||
baseline=baseline,
|
||||
target=target,
|
||||
name="extreme_values",
|
||||
tmp_path=tmp_path,
|
||||
publish_dir=publish_dir,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
96
test/registered/debug_utils/comparator/test_visualizer.py
Normal file
96
test/registered/debug_utils/comparator/test_visualizer.py
Normal file
@@ -0,0 +1,96 @@
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.debug_utils.comparator.visualizer.preprocessing import (
|
||||
_preprocess_tensor,
|
||||
_reshape_to_balanced_aspect,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=30, suite="default", nightly=True)
|
||||
|
||||
|
||||
class TestPreprocessTensor:
|
||||
def test_1d_becomes_2d(self) -> None:
|
||||
t: torch.Tensor = torch.randn(100)
|
||||
result: torch.Tensor = _preprocess_tensor(t)
|
||||
assert result.ndim == 2
|
||||
|
||||
def test_3d_becomes_2d(self) -> None:
|
||||
t: torch.Tensor = torch.randn(2, 3, 4)
|
||||
result: torch.Tensor = _preprocess_tensor(t)
|
||||
assert result.ndim == 2
|
||||
assert result.numel() == t.numel()
|
||||
|
||||
def test_high_dim_becomes_2d(self) -> None:
|
||||
t: torch.Tensor = torch.randn(2, 3, 4, 5)
|
||||
result: torch.Tensor = _preprocess_tensor(t)
|
||||
assert result.ndim == 2
|
||||
assert result.numel() == t.numel()
|
||||
|
||||
def test_scalar_becomes_2d(self) -> None:
|
||||
t: torch.Tensor = torch.tensor(3.14)
|
||||
result: torch.Tensor = _preprocess_tensor(t)
|
||||
assert result.ndim == 2
|
||||
assert result.numel() == 1
|
||||
|
||||
def test_already_2d_preserves_elements(self) -> None:
|
||||
t: torch.Tensor = torch.randn(10, 20)
|
||||
result: torch.Tensor = _preprocess_tensor(t)
|
||||
assert result.ndim == 2
|
||||
assert result.numel() == 200
|
||||
|
||||
|
||||
class TestReshapeToBalancedAspect:
|
||||
def test_extreme_wide_gets_fixed(self) -> None:
|
||||
t: torch.Tensor = torch.randn(1, 10000)
|
||||
result: torch.Tensor = _reshape_to_balanced_aspect(t)
|
||||
h, w = result.shape
|
||||
ratio: float = max(h, w) / max(min(h, w), 1)
|
||||
assert ratio <= 5.0
|
||||
|
||||
def test_extreme_tall_gets_fixed(self) -> None:
|
||||
t: torch.Tensor = torch.randn(10000, 1)
|
||||
result: torch.Tensor = _reshape_to_balanced_aspect(t)
|
||||
h, w = result.shape
|
||||
ratio: float = max(h, w) / max(min(h, w), 1)
|
||||
assert ratio <= 5.0
|
||||
|
||||
def test_already_balanced_unchanged(self) -> None:
|
||||
t: torch.Tensor = torch.randn(100, 100)
|
||||
result: torch.Tensor = _reshape_to_balanced_aspect(t)
|
||||
assert result.shape == (100, 100)
|
||||
|
||||
def test_preserves_numel(self) -> None:
|
||||
t: torch.Tensor = torch.randn(1, 7919)
|
||||
result: torch.Tensor = _reshape_to_balanced_aspect(t)
|
||||
assert result.numel() == t.numel()
|
||||
|
||||
|
||||
class TestGenerateComparisonFigure:
|
||||
@pytest.fixture(autouse=True)
|
||||
def _skip_if_no_matplotlib(self) -> None:
|
||||
pytest.importorskip("matplotlib")
|
||||
|
||||
def test_nested_output_dir(self, tmp_path: Path) -> None:
|
||||
from sglang.srt.debug_utils.comparator.visualizer import (
|
||||
generate_comparison_figure,
|
||||
)
|
||||
|
||||
output_path: Path = tmp_path / "a" / "b" / "c" / "nested.png"
|
||||
|
||||
generate_comparison_figure(
|
||||
baseline=torch.randn(10, 10),
|
||||
target=torch.randn(10, 10),
|
||||
name="nested",
|
||||
output_path=output_path,
|
||||
)
|
||||
|
||||
assert output_path.exists()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
Reference in New Issue
Block a user