test-writing

name

test-writing

description

Guide for writing tests in torch-harmonics. Use this skill whenever the user asks to add tests, check or extend test coverage, write a new test class, test a new layer or kernel, or add distributed tests. Also use when the user asks about tolerance values, how to compare tensors, how to structure distributed test infrastructure, or how to test gradients. Covers serial tests in test_convolution.py and distributed tests in test_distributed_*.py.

Test Writing Guide

General principles

• Match existing style exactly. Before writing any test, read the relevant existing test class in tests/ to pick up naming conventions, helper usage, and parameterized.expand patterns.
• Numerical tests against a reference baseline. Every forward and backward check must compare against a deterministic reference (typically fp32 serial or the torch reference path). Never just check shapes or dtypes alone.
• Test every code path for every precision. If a layer has 4 dispatch branches (e.g. fused×kpacked), run each one. If it supports fp32/fp16/bf16, run all three.

---

Naming Conventions

The files for serial tests should be named test_<feature_name>.py and distributed tests test_distributed_<feature_name>.py respectively.

---

Default tolerances

dtype	atol	rtol
fp64	1e-6	1e-6
fp32	1e-6	1e-5
fp16	1e-3	1e-2
bf16	1e-2	5e-2

Parameter gradients accumulate noise over batch × spatial dims, so loosen their tolerance by a factor of 10–1000× relative to per-element output tolerances (use judgment based on reduction size). For stacked/chained layers, loosen further (2–5×).

---

Comparison helpers

Tensor comparisons — always use compare_tensors. The first argument is a name string used in failure messages; it is required:

ok = compare_tensors("output", ref, got, atol=atol, rtol=rtol, verbose=verbose)
self.assertTrue(ok, "output")

Python logic checks — use plain assertTrue / assertFalse:

self.assertIsNotNone(conv.psi_kpacked_K_pad)
self.assertIn(conv.psi_kpacked_K_pad, (8, 16))
self.assertEqual(out.dtype, torch.bfloat16)

Never use torch.allclose or torch.testing.assert_close directly — these don't give the diagnostic output that compare_tensors provides.

verbose flag

Every test method that calls compare_tensors must accept a verbose=False parameter and forward it:

def test_my_feature(self, ..., verbose=False):
    ok = compare_tensors("output", ref, got, atol=atol, rtol=rtol, verbose=verbose)
    self.assertTrue(ok)

In distributed tests, only rank 0 should print:

verbose = verbose and self.world_rank == 0

---

Autocast / AMP

Wrap forward (and any backward that needs grad under the same dtype) with the maybe_autocast helper from testutils:

with maybe_autocast(self.device.type, dtype):
    out = module(inp)

For AMP dtypes (float16, bfloat16), keep module parameters and inputs in float32 and let autocast downcast inside the forward. Match this pattern exactly — do not manually cast module weights or inputs before the call.

is_amp = dtype in (torch.float16, torch.bfloat16)
module_dtype = torch.float32 if is_amp else dtype
module = MyLayer(...).to(dtype=module_dtype, device=device)
inp = torch.randn(..., dtype=module_dtype, device=device)
with maybe_autocast(device.type, dtype):
    out = module(inp)

---

What to test in every test method

Test all three of these in a single test method; do not split them across separate tests unless there is a strong reason:

1. Forward output — compare against reference with compare_tensors.
2. Input gradient — run .backward(ograd), compare inp.grad with compare_tensors.
3. Weight (parameter) gradients — compare each parameter's .grad with compare_tensors. See distributed section for the all-reduce needed there.

inp.requires_grad_(True)
with maybe_autocast(device.type, dtype):
    out = module(inp)
ograd = torch.randn_like(out)
out.backward(ograd)

ok = compare_tensors("output",   ref_out,   out,      atol=atol, rtol=rtol, verbose=verbose)
self.assertTrue(ok, "output")
ok = compare_tensors("inp grad", ref_igrad, inp.grad, atol=atol, rtol=rtol, verbose=verbose)
self.assertTrue(ok, "inp grad")
if module.weight.grad is not None:
    ok = compare_tensors("weight grad", ref_wgrad, module.weight.grad,
                         atol=wgrad_atol, rtol=wgrad_rtol, verbose=verbose)
    self.assertTrue(ok, "weight grad")

---

Serial test structure

class TestMyFeature(unittest.TestCase):
    device = torch.device("cuda")  # or "cpu"

    @classmethod
    def setUpClass(cls):
        disable_tf32()

    def test_my_case(self, ..., verbose=False):
        set_seed(42)
        ...

disable_tf32() must be called in setUpClass for any test class that runs on CUDA — TF32 reduces fp32 matmul precision and causes spurious failures at the default tolerances.

For parameterised cases use @parameterized.expand([...], skip_on_empty=True):

@parameterized.expand([
    # [param1, param2, ..., atol, rtol]
    [64, 128, torch.float32, 1e-6, 1e-5],
    [64, 128, torch.bfloat16, 1e-2, 5e-2],
], skip_on_empty=True)
def test_my_feature(self, nlat, nlon, dtype, atol, rtol, verbose=False):
    ...

---

Distributed test structure

Module-level wireup (required in every distributed test file)

_DIST_CTX = {}

def setUpModule():
    setup_module(_DIST_CTX)

def tearDownModule():
    teardown_module(_DIST_CTX)

Class setup

class TestMyDistributedFeature(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        setup_class_from_context(cls, _DIST_CTX)
        disable_tf32()

    def setUp(self):
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        if dist.is_initialized():
            dist.barrier()

    def tearDown(self):
        if torch.cuda.is_available():
            torch.cuda.synchronize()

setup_class_from_context sets cls.device, cls.world_rank, cls.grid_size_h, cls.grid_size_w, cls.hrank, cls.wrank, cls.h_group, cls.w_group on the class.

disable_tf32() is called in setUpClass for the same reason as serial tests.

Comparing distributed output against serial baseline

The pattern is always:

1. Build a serial (local) module and a distributed module with identical weights.
2. Run serial forward + backward on the full tensor.
3. Split the input, run distributed forward + backward on the local shard.
4. Gather distributed outputs/grads back to the full tensor.
5. Compare with compare_tensors on all ranks; use reduce_success so a failure on any rank fails the test.

# 1. build modules
conv_local = MyLayer(**args).to(dtype=module_dtype, device=self.device)
conv_dist  = DistributedMyLayer(**args).to(dtype=module_dtype, device=self.device)

# 2. sync weights — every rank must have identical parameters
with torch.no_grad():
    conv_dist.weight.copy_(conv_local.weight)
    if hasattr(conv_dist, "bias") and conv_dist.bias is not None:
        conv_dist.bias.copy_(conv_local.bias)

# 3. serial fwd + bwd
inp_full = torch.randn((B, C, H, W), dtype=module_dtype, device=self.device)
inp_full.requires_grad_(True)
with maybe_autocast(self.device.type, dtype):
    out_full = conv_local(inp_full)
ograd_full = torch.randn_like(out_full)
out_full.backward(ograd_full)
igrad_full = inp_full.grad.clone()

# 4. distributed fwd + bwd on local shard
inp_local = self._split_helper(inp_full.detach().clone())
inp_local.requires_grad_(True)
with maybe_autocast(self.device.type, dtype):
    out_local = conv_dist(inp_local)
ograd_local = self._split_helper(ograd_full)
out_local.backward(ograd_local)
igrad_local = inp_local.grad.clone()

Gather helpers

Use the gather_tensor_hw utility (from testutils) to gather distributed tensors back to the full shape on all ranks:

def _split_helper(self, tensor):
    return split_tensor_hw(tensor,
        hdim=-2, wdim=-1,
        hsize=self.grid_size_h, wsize=self.grid_size_w,
        hrank=self.hrank, wrank=self.wrank)

def _gather_helper_fwd(self, tensor, conv_dist):
    return gather_tensor_hw(tensor,
        hdim=-2, wdim=-1,
        hshapes=conv_dist.lat_out_shapes, wshapes=conv_dist.lon_out_shapes,
        hsize=self.grid_size_h, wsize=self.grid_size_w,
        hrank=self.hrank, wrank=self.wrank,
        hgroup=self.h_group, wgroup=self.w_group)

def _gather_helper_bwd(self, tensor, conv_dist):
    return gather_tensor_hw(tensor,
        hdim=-2, wdim=-1,
        hshapes=conv_dist.lat_in_shapes, wshapes=conv_dist.lon_in_shapes,
        hsize=self.grid_size_h, wsize=self.grid_size_w,
        hrank=self.hrank, wrank=self.wrank,
        hgroup=self.h_group, wgroup=self.w_group)

Use _gather_helper_fwd for output tensors (indexed by output spatial shape) and _gather_helper_bwd for input-space tensors (input gradient).

Comparing gathered results

Always use reduce_success so a test failure on any rank propagates to all:

verbose = verbose and self.world_rank == 0

out_gather = self._gather_helper_fwd(out_local, conv_dist)
ok = compare_tensors("output", out_full, out_gather, atol=atol, rtol=rtol, verbose=verbose)
self.assertTrue(reduce_success(ok, self.device), "output")

igrad_gather = self._gather_helper_bwd(igrad_local, conv_dist)
ok = compare_tensors("gradients", igrad_full, igrad_gather, atol=atol, rtol=rtol, verbose=verbose)
self.assertTrue(reduce_success(ok, self.device), "gradients")

Weight gradient reductions for distributed layers

Distributed layers compute parameter gradients as partial sums over the local spatial tile. To compare against the serial reference, all-reduce across the polar (h) and azimuth (w) groups independently:

def _allreduce_param_grad(self, tensor):
    out = tensor.clone()
    if self.grid_size_h > 1:
        dist.all_reduce(out, group=self.h_group)
    if self.grid_size_w > 1:
        dist.all_reduce(out, group=self.w_group)
    return out

Then compare:

param_tol = 1000.0 if not is_amp else 10.0   # param grads accumulate more noise
pg_atol, pg_rtol = atol * param_tol, rtol * param_tol
if conv_dist.weight.grad is not None:
    wgrad = self._allreduce_param_grad(conv_dist.weight.grad)
    ok = compare_tensors("weight grad", conv_local.weight.grad, wgrad,
                         atol=pg_atol, rtol=pg_rtol, verbose=verbose)
    self.assertTrue(reduce_success(ok, self.device), "weight grad")

The tighter tolerance factor for AMP is intentional: the base atol/rtol for fp16/bf16 is already 3–4 orders of magnitude looser than fp32, so multiplying by 1000× makes the bound meaninglessly large.

---

Checklist before committing a new test

• disable_tf32() called in setUpClass
• set_seed(N) called at the top of every test method that uses random tensors
• verbose=False parameter on every method that calls compare_tensors
• All three of forward output, input gradient, weight gradients are tested
• All relevant precisions are covered (fp32, fp16, bf16 at minimum for CUDA layers)
• All dispatch branches are covered (e.g. fused=True/False, kpacked-enabled/disabled)
• Distributed tests use setUpModule / tearDownModule wireup
• Distributed weight grads are all-reduced before comparison
• reduce_success used for every assertion in distributed tests
• AMP tests use maybe_autocast — not manual .to(dtype) on module weights