From 5e0b50c1518bac220ebbe0a60c24d23520150342 Mon Sep 17 00:00:00 2001 From: x66ccff Date: Tue, 6 Oct 2026 21:56:35 +0800 Subject: [PATCH] Reduce PSRN scoring memory with ordered candidate tiles --- psrn/model/models.py | 303 ++++++++++++++++++++++++-- psrn/model/regressor.py | 26 +-- test/test_dense_inplace_scoring.py | 330 +++++++++++++++++++++++++++++ test/test_row_aligned_tiles.py | 190 +++++++++++++++++ test/test_scoring_dispatch.py | 166 +++++++++++++++ 5 files changed, 979 insertions(+), 36 deletions(-) create mode 100644 test/test_dense_inplace_scoring.py create mode 100644 test/test_row_aligned_tiles.py create mode 100644 test/test_scoring_dispatch.py diff --git a/psrn/model/models.py b/psrn/model/models.py index 669eb0e..3c1976e 100644 --- a/psrn/model/models.py +++ b/psrn/model/models.py @@ -1,4 +1,5 @@ -import torch +import math +import torch import torch.nn as nn from .functions import ( @@ -18,6 +19,14 @@ ) from .functions import Sign, Pow2, Pow3, Pow, Sigmoid, Abs, Cosh, Tanh, Sqrt + + +def _dtype_element_size(dtype): + try: + return dtype.itemsize + except AttributeError: + # Older Torch exposes size on tensors rather than the dtype object. + return torch.empty((), dtype=dtype, device="cpu").element_size() # Duplicate Removal Layer @@ -108,8 +117,106 @@ def forward(self, x): h = torch.cat(h, dim=1) # shape: (batch_size, out_dim) return h - def init_offset(self, device): - self.offset_tensor = self.get_offset_tensor(device) + def iter_tiles(self, x, tile_size): + """Yield (global offset, values) without a full final-layer output. + + Inputs may contain multiple independent rows. Candidate order is the + same operator-block/operand order as ``forward`` for a single row. + """ + if not isinstance(tile_size, int) or isinstance(tile_size, bool) or tile_size <= 0: + raise ValueError("tile_size must be a positive integer") + if x.dim() != 2 or x.shape[1] != self.in_dim: + raise ValueError("tile inputs must have shape (samples, in_dim)") + binary_operations = { + Add: torch.add, Mul: torch.mul, + Sub: torch.sub, SemiSub: torch.sub, + Div: torch.div, SemiDiv: torch.div, Pow: torch.pow, + } + triangle_starts = None + if self.n_binary_U and self.in_dim > tile_size: + rows = torch.arange(self.in_dim + 1, device=x.device, dtype=torch.int64) + triangle_starts = rows * (2 * self.in_dim - rows + 1) // 2 + offset = 0 + for module in self.list: + count = module.out_dim + if module.is_unary: + # Preserve the original single-row/full-width kernel shape. + # CPU transcendental vector/tail kernels can round differently + # if rows are batched or the unary input itself is tiled. + unary_values = torch.cat( + [module(x[row:row + 1]) for row in range(x.shape[0])], dim=0 + ) + elif self.in_dim <= tile_size: + operation = binary_operations.get(type(module)) + if operation is None: + raise NotImplementedError("tiled scoring for {}".format(type(module).__name__)) + row_start = 0 + while row_start < self.in_dim: + if module.is_directed: + row_end = min(self.in_dim, row_start + tile_size // self.in_dim) + start = row_start * self.in_dim + # Views of complete rows preserve the original directed + # broadcast geometry without candidate-sized gathers. + values = operation( + x[:, row_start:row_end, None], x[:, None, :] + ).reshape(x.shape[0], -1) + else: + start = row_start * (2 * self.in_dim - row_start + 1) // 2 + # Find the last complete triangular row within the cap. + # All arithmetic stays on the host and uses exact ints. + low, high = row_start + 1, self.in_dim + 1 + while high - low > 1: + middle = (low + high) // 2 + stop = middle * (2 * self.in_dim - middle + 1) // 2 + if stop - start <= tile_size: + low = middle + else: + high = middle + row_end = low + indices = torch.triu_indices( + row_end - row_start, self.in_dim, offset=row_start, + dtype=torch.int64, device=x.device, + ) + indices[0].add_(row_start) + values = operation(x[:, indices[0]], x[:, indices[1]]) + yield offset + start, values + del values + if not module.is_directed: + del indices + row_start = row_end + offset += count + continue + for start in range(0, count, tile_size): + end = min(start + tile_size, count) + if module.is_unary: + values = unary_values[:, start:end] + else: + operation = binary_operations.get(type(module)) + if operation is None: + raise NotImplementedError("tiled scoring for {}".format(type(module).__name__)) + local = torch.arange(start, end, device=x.device, dtype=torch.int64) + if module.is_directed: + left = torch.div(local, self.in_dim, rounding_mode="floor") + right = local % self.in_dim + else: + left = torch.searchsorted(triangle_starts, local, right=True) - 1 + right = left + local - triangle_starts[left] + # Match functions.forward, including unprotected division, + # powers, and their NaN/Inf behavior. + values = operation(x[:, left], x[:, right]) + yield offset + start, values + del values + if not module.is_unary: + del local, left, right + if module.is_unary: + del unary_values + offset += count + + def init_offset(self, device): + # Operand indices are reconstructed only for selected expressions. + # Keep the dense-table builder available for compatibility and audits, + # but do not materialize its candidate-sized storage at initialization. + self.offset_tensor = None def get_offset_tensor(self, device): offset_tensor = torch.zeros((self.out_dim, 2), dtype=torch.int, device=device) @@ -170,15 +277,31 @@ def get_out_dim_cum_ls(self): self.out_dim_cum_ls = [sum(out_dim_ls[: i + 1]) for i in range(len(out_dim_ls))] return self.out_dim_cum_ls - def get_op_and_offset(self, index): - out_dim_cum_ls = self.get_out_dim_cum_ls() - for i, func in enumerate(self.list): - if index < out_dim_cum_ls[i]: - break - offset = self.offset_tensor[ - index - ].tolist() # index 65 is out of bounds for dimension 0 with size 65 - return func.operator, offset + def get_op_and_offset(self, index): + if index < 0 or index >= self.out_dim: + raise IndexError("candidate index out of range") + out_dim_cum_ls = self.get_out_dim_cum_ls() + start = 0 + for i, func in enumerate(self.list): + if index < out_dim_cum_ls[i]: + break + start = out_dim_cum_ls[i] + local = index - start + if func.is_unary: + offset = [local, self.in_dim] + elif func.is_directed: + offset = [local // self.in_dim, local % self.in_dim] + else: + # Invert row-major upper-triangle enumeration using integer + # arithmetic. The initial root may lie one row past the answer. + d = self.in_dim + row = (2 * d + 1 - math.isqrt((2 * d + 1) ** 2 - 8 * local)) // 2 + prefix = row * (2 * d - row + 1) // 2 + if prefix > local: + row -= 1 + prefix = row * (2 * d - row + 1) // 2 + offset = [row, row + local - prefix] + return func.operator, offset class PSRN(nn.Module): @@ -255,6 +378,162 @@ def forward(self, x): h = layer(h) return h # shape: (batch_size, out_dim) + def estimate_dense_score_bytes(self, x, y): + """Conservative live-tensor bound for scoring one dense prediction row. + + Count the score buffer, layer input/output, operator outputs retained + until cat, triangular int32 operand indices and two gathered operands, + and the prediction/error/square at accumulation. This is a dispatch + estimate, not an allocator/reserved-memory or whole-process limit. + """ + value_bytes = _dtype_element_size(x.dtype) + # A zero-dimensional target has scalar promotion rules, unlike a + # one-element vector. No tensor values or device state are read. + error_dtype = torch.result_type(x, y[0]) + error_bytes = _dtype_element_size(error_dtype) + score_bytes = self.out_dim * _dtype_element_size(torch.get_default_dtype()) + peak = 0 + for layer in self.list: + inputs = layer.in_dim * value_bytes + outputs = layer.out_dim * value_bytes + # All operator results coexist with the concatenated output. + layer_peak = inputs + 2 * outputs + if isinstance(layer, SymbolLayer): + largest_triangle = 0 + for op in layer.list: + if not op.is_unary and not op.is_directed: + largest_triangle = max(largest_triangle, op.out_dim) + # Only the largest triangular block can set this bound. + # Inspect current modules rather than cache mutable topology. + if largest_triangle: + layer_peak = max( + layer_peak, + inputs + outputs + largest_triangle * (8 + 2 * value_bytes), + ) + peak = max(peak, layer_peak) + peak = max(peak, self.out_dim * (value_bytes + 2 * error_bytes)) + return score_bytes + peak + + @torch.no_grad() + def score_mse(self, x, y, tile_size=1048576, sample_chunk_size=16, + strategy="auto", workspace_budget_bytes=512 * 1024 ** 2): + """Return ordered candidate MSEs using dense rows or final-layer tiles. + + Auto uses the original single-row forward when its estimated live + tensor storage fits a sample-count-aware allowance. The allowance is + workspace_budget_bytes * sample_chunk_size / average_chunk_rows, + where average_chunk_rows counts actual rows per tiled sample group. + Thus small queries can stay dense when tiling has little batching + benefit. The budget is a policy base, not a hard workspace cap; the + global score vector is always retained. Use strategy="tiled" to + exercise tiles explicitly, or "dense" for a reference evaluation. + Both paths keep original accumulator dtype, sample order and global + min/topk ties. CPU Pow, active autocast and nonfloating inputs always + use dense rows to preserve rounding, final cat dtype promotion and + the original behavior for unsupported input types. + last_score_strategy records the selected path without device queries. + """ + if x.dim() != 2 or x.shape[0] == 0 or x.shape[1] != self.n_variables: + raise ValueError("x must contain at least one row of n_variables inputs") + sample_count = x.shape[0] + target_rank = y.dim() + if target_rank == 0 or y.shape[0] != sample_count or y.numel() != sample_count: + raise ValueError("y must contain one scalar target per input row") + for name, value in (("tile_size", tile_size), ("sample_chunk_size", sample_chunk_size), + ("workspace_budget_bytes", workspace_budget_bytes)): + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError("{} must be a positive integer".format(name)) + if strategy not in ("auto", "dense", "tiled"): + raise ValueError("strategy must be auto, dense, or tiled") + # Preserve scalar targets for (N,), and vector promotion for every + # higher rank, without checking or reshaping each target row. + if target_rank > 2: + y = y.reshape(sample_count, 1) + # AMP may promote operator outputs before the original final cat. + # Preserve that shared prediction dtype as well as AMP pow semantics. + try: + autocast_enabled = torch.is_autocast_enabled(x.device.type) + except TypeError: + # Torch before 2.4 has separate no-argument CUDA/CPU queries. + if x.device.type == "cuda": + autocast_enabled = torch.is_autocast_enabled() + elif x.device.type == "cpu": + autocast_enabled = torch.is_autocast_cpu_enabled() + else: + # An unknown legacy backend keeps the original square. + autocast_enabled = True + cpu_pow = x.device.type == "cpu" and any( + isinstance(op, Pow) for op in self.list[-1].list + ) + # Integer operators can retain integer outputs while transcendental + # operators promote to float. Score only after the shared final cat. + # Complex inputs likewise retain original forward/scoring behavior. + use_dense = (autocast_enabled or not x.is_floating_point() + or cpu_pow or strategy == "dense") + if not use_dense and strategy == "auto": + chunk_rows = min(sample_chunk_size, sample_count) + sample_groups = (sample_count + chunk_rows - 1) // chunk_rows + # Compare without floating point rounding at the policy boundary. + use_dense = self.estimate_dense_score_bytes(x, y) * sample_count <= ( + workspace_budget_bytes * sample_chunk_size * sample_groups + ) + self.last_score_strategy = "dense" if use_dense else "tiled" + if use_dense: + sums = torch.zeros((1, self.out_dim), device=self.device) + for sample in range(sample_count): + h = self.forward(x[sample].reshape(1, -1)) + target = y[sample] + diff = h - target + if autocast_enabled: + sums += diff ** 2 + else: + diff.pow_(2) + sums += diff + # Do not carry the previous full row into the next forward. + del h, diff + sums /= sample_count + scores = sums.reshape(-1) + scores.nan_to_num_(nan=float("inf"), posinf=float("inf"), neginf=float("inf")) + return scores + scores = torch.zeros(self.out_dim, device=self.device) + final_layer = self.list[-1] + # Each sample chunk completes all candidate tiles before the next + # chunk; each candidate still accumulates samples in original order. + for sample_start in range(0, sample_count, sample_chunk_size): + sample_end = min(sample_start + sample_chunk_size, sample_count) + # Retain only this chunk's pre-final features. Directed earlier + # layers expect a single row, so they must not receive a batch. + if len(self.list) == 1: + # Preserve original row strides for unary kernel selection. + features = x[sample_start:sample_end] + else: + features = [] + for sample in range(sample_start, sample_end): + h = x[sample].reshape(1, -1) + for layer in self.list[:-1]: + h = layer(h) + features.append(h) + features = torch.cat(features, dim=0) + for offset, predictions in final_layer.iter_tiles(features, tile_size): + sums = scores[offset:offset + predictions.shape[1]] + for local_sample in range(sample_end - sample_start): + # Preserve the original target shape and dtype promotion: + # y[i] is scalar for (N,), one-element for (N, 1). + target = y[sample_start + local_sample] + diff = predictions[local_sample] - target + square = diff ** 2 + sums += square + # Release consumer references before the generator allocates + # the next tile; its own references are cleared after yield. + del predictions, diff, square + # In-place division avoids a second full candidate vector. Element + # arithmetic and normalization order are identical to the baseline. + scores /= sample_count + # Preserve every finite value and map all nonfinite scores to +Inf, + # without candidate-sized abs and boolean-indexing intermediates. + scores.nan_to_num_(nan=float("inf"), posinf=float("inf"), neginf=float("inf")) + return scores + def get_expr(self, index): return self._get_expr(index, -1) diff --git a/psrn/model/regressor.py b/psrn/model/regressor.py index eba833c..9954b8b 100644 --- a/psrn/model/regressor.py +++ b/psrn/model/regressor.py @@ -1427,18 +1427,7 @@ def get_reward(self, eta, complexity, mse): def get_best_expr_and_MSE(self, X, Y): with torch.no_grad(): - sum_ = torch.zeros((1, self.net.out_dim), device=self.net.device) - for i in range(X.shape[0]): - H = self.net.forward(X[i].reshape(1, -1)) - diff = H - Y[i] - square = diff**2 - sum_ += square - mean = sum_ / X.shape[0] - mean = mean.reshape(-1) - - # replace all nan, -inf to inf - mean[torch.isnan(mean)] = float("inf") - mean[torch.isinf(mean)] = float("inf") + mean = self.net.score_mse(X, Y) min_value, min_index = torch.min(mean, dim=0) @@ -1454,18 +1443,7 @@ def get_best_expr_and_MSE_topk(self, X, Y, n_top): self.fitted_expr_c_set = set() with torch.no_grad(): - sum_ = torch.zeros((1, self.net.out_dim), device=self.net.device) - for i in range(X.shape[0]): - H = self.net.forward(X[i].reshape(1, -1)) - diff = H - Y[i] - square = diff**2 - sum_ += square - mean = sum_ / X.shape[0] - mean = mean.reshape(-1) - - # replace all nan, -inf to inf - mean[torch.isnan(mean)] = float("inf") - mean[torch.isinf(mean)] = float("inf") + mean = self.net.score_mse(X, Y) values, indices = torch.topk(mean, n_top, largest=False, sorted=True) MSE_min_ls = values.tolist() diff --git a/test/test_dense_inplace_scoring.py b/test/test_dense_inplace_scoring.py new file mode 100644 index 0000000..8a48a8d --- /dev/null +++ b/test/test_dense_inplace_scoring.py @@ -0,0 +1,330 @@ +"""Check dense scratch reduction against the unchanged ordered MSE formula.""" +import importlib +import itertools +from pathlib import Path +import sys +import types +import unittest +from unittest.mock import patch + +import torch + +package = types.ModuleType('_psrn_dense_inplace_tests') +package.__path__ = [str(Path(__file__).resolve().parents[1] / 'psrn')] +sys.modules[package.__name__] = package +models = importlib.import_module(package.__name__ + '.model.models') +PSRN = models.PSRN +FAMILIES = [ + ['Add', 'Mul', 'Sub', 'Div', 'Identity', 'Sin', 'Cos', 'Exp', 'Log'], + ['Add', 'Mul', 'SemiSub', 'SemiDiv', 'Identity', 'Sin', 'Cos', 'Exp', 'Log'], + ['Add', 'Mul', 'Identity', 'Pow2', 'Pow3', 'Neg'], + ['Pow', 'Identity'], + ['Identity', 'Sin', 'Cos', 'Exp', 'Log', 'Neg', 'Inv', 'Sign', 'Pow2', + 'Pow3', 'Sigmoid', 'Abs', 'Cosh', 'Tanh', 'Sqrt'], + ['Add', 'Mul', 'Sub', 'Div', 'SemiSub', 'SemiDiv', 'Pow', 'Identity', + 'Sin', 'Cos', 'Exp', 'Log', 'Neg', 'Inv', 'Sign', 'Pow2', 'Pow3', + 'Sigmoid', 'Abs', 'Cosh', 'Tanh', 'Sqrt'], +] + + +@torch.no_grad() +def original_scores(net, x, y): + if y.dim() > 2: + y = y.reshape(x.shape[0], 1) + sums = torch.zeros((1, net.out_dim), device=net.device) + for row in range(x.shape[0]): + h = net.forward(x[row].reshape(1, -1)) + diff = h - y[row] + square = diff ** 2 + sums += square + result = (sums / x.shape[0]).reshape(-1) + result[~torch.isfinite(result)] = float('inf') + return result + + +def integer_view(values): + return values.contiguous().view(torch.int64 if values.element_size() == 8 else torch.int32) + + +class DenseInplaceScoringTests(unittest.TestCase): + @classmethod + def setUpClass(cls): + torch.set_num_threads(1) + + def test_full_scores_ties_dtype_promotion_masks_and_strides(self): + for operators, masked in itertools.product(FAMILIES, (False, True)): + prefix = PSRN(2, operators, 1, device='cpu') + mask = torch.arange(prefix.out_dim) % 3 != 1 if masked else None + net = PSRN(2, operators, 2, dr_mask=mask, device='cpu') + for xd, yd, rank, samples, invalid in itertools.product( + (torch.float32, torch.float64), (torch.float32, torch.float64), + (1, 2, 4), (1, 10, 17), (False, True)): + with self.subTest(operators=operators, masked=masked, xd=xd, yd=yd, + rank=rank, samples=samples, invalid=invalid): + storage = torch.linspace(.1, 1.7, samples * 8, dtype=xd).reshape(samples * 2, 4) + x = storage[::2, ::2] + if invalid: + x[0, 0] = 0 + x[0, 1] = -1 + if samples > 1: + x[1, 0], x[2, 0] = float('nan'), float('inf') + x[3, 0], x[4, 0] = -float('inf'), 10000 + y = torch.linspace(-1, 1, samples * 2, dtype=yd)[::2] + y = y.reshape((samples,) + (1,) * (rank - 1)) + x_bits = integer_view(x).clone() + y_bits = integer_view(y).clone() + expected = original_scores(net, x, y) + actual = net.score_mse(x, y, strategy='dense') + self.assertTrue(torch.equal(integer_view(actual), integer_view(expected))) + k = min(5, net.out_dim) + self.assertTrue(torch.equal(torch.topk(actual, k, largest=False).indices, + torch.topk(expected, k, largest=False).indices)) + self.assertTrue(torch.equal(integer_view(x), x_bits)) + self.assertTrue(torch.equal(integer_view(y), y_bits)) + self.assertEqual(net.last_score_strategy, 'dense') + + def test_autocast_keeps_original_pow_path(self): + net = PSRN(2, ['Identity', 'Add', 'Pow2'], 1, device='cpu') + real_pow = torch.Tensor.pow_ + calls = [] + + def observed_pow(values, exponent): + calls.append(values.shape) + return real_pow(values, exponent) + + for dtype in (torch.bfloat16, torch.float32, torch.float64): + x = torch.tensor([[.25, .5], [1., -.75], [.125, 2.]], dtype=dtype) + y = torch.tensor([.125, -1., .5], dtype=dtype) + with torch.autocast('cpu', dtype=torch.bfloat16): + expected = original_scores(net, x, y) + with patch.object(torch.Tensor, 'pow_', new=observed_pow): + actual = net.score_mse(x, y, strategy='dense') + self.assertEqual(calls, []) + self.assertTrue(torch.equal(integer_view(actual), integer_view(expected))) + with patch.object(torch.Tensor, 'pow_', new=observed_pow): + actual = net.score_mse(x, y, strategy='dense') + self.assertEqual(len(calls), x.shape[0]) + self.assertTrue(torch.equal(integer_view(actual), integer_view(original_scores(net, x, y)))) + + def test_legacy_autocast_query_arity_cpu_cuda_and_unknown_metadata(self): + # All computation stays on CPU. Only device-query metadata is mocked + # to exercise CUDA and unknown legacy API selection without hardware. + net = PSRN(2, ['Identity'], 1, device='cpu') + x = torch.tensor([[.25, -.5], [1.25, 2.]]) + y = torch.tensor([.1, -.2]) + expected = original_scores(net, x, y) + real_pow = torch.Tensor.pow_ + for device, enabled in (('cpu', False), ('cpu', True), + ('cuda', False), ('cuda', True), ('mps', False)): + calls = [] + + def legacy_query(*args): + if args: + raise TypeError('legacy query accepts no arguments') + return enabled + + def observed_pow(values, exponent): + calls.append(values.shape) + return real_pow(values, exponent) + + with self.subTest(device=device, enabled=enabled): + with patch.object(torch, 'is_autocast_enabled', side_effect=legacy_query) as query, \ + patch.object(torch, 'is_autocast_cpu_enabled', return_value=enabled) as cpu_query, \ + patch.object(torch.Tensor, 'device', new=property(lambda _: torch.device(device))), \ + patch.object(torch.Tensor, 'pow_', new=observed_pow): + actual = net.score_mse(x, y, strategy='dense') + self.assertEqual(query.call_args_list[0].args, (device,)) + if device == 'cpu': + self.assertEqual(query.call_count, 1) + cpu_query.assert_called_once_with() + elif device == 'cuda': + self.assertEqual(query.call_count, 2) + self.assertEqual(query.call_args_list[1].args, ()) + cpu_query.assert_not_called() + else: + self.assertEqual(query.call_count, 1) + cpu_query.assert_not_called() + expected_pow_calls = x.shape[0] if device != 'mps' and not enabled else 0 + self.assertEqual(len(calls), expected_pow_calls) + self.assertTrue(torch.equal(integer_view(actual), integer_view(expected))) + + def test_modern_autocast_query_avoids_legacy_fallback(self): + net = PSRN(2, ['Identity'], 1, device='cpu') + x, y = torch.tensor([[.1, .5]]), torch.tensor([.25]) + with patch.object(torch, 'is_autocast_enabled', wraps=torch.is_autocast_enabled) as query, \ + patch.object(torch, 'is_autocast_cpu_enabled', side_effect=AssertionError('legacy called')): + actual = net.score_mse(x, y, strategy='dense') + query.assert_called_once_with('cpu') + self.assertTrue(torch.equal(integer_view(actual), integer_view(original_scores(net, x, y)))) + + def test_dtype_size_modern_metadata_needs_no_allocation(self): + with patch.object(torch, 'empty', side_effect=AssertionError('modern allocation')): + for dtype in (torch.float32, torch.float64, torch.bfloat16): + self.assertEqual(models._dtype_element_size(dtype), dtype.itemsize) + + def test_legacy_dtype_size_auto_bound_route_and_scores(self): + class LegacyDtype: + def __init__(self, actual): + self.actual = actual + + real_empty = torch.empty + real_result_type = torch.result_type + real_default_dtype = torch.get_default_dtype + real_tensor_dtype = torch.Tensor.dtype + allocations = [] + + def legacy_empty(*args, **kwargs): + self.assertEqual(args, ((),)) + self.assertEqual(kwargs['device'], 'cpu') + self.assertIsInstance(kwargs['dtype'], LegacyDtype) + allocations.append(kwargs['dtype'].actual) + kwargs['dtype'] = kwargs['dtype'].actual + return real_empty(*args, **kwargs) + + def legacy_tensor_dtype(tensor): + return LegacyDtype(real_tensor_dtype.__get__(tensor, torch.Tensor)) + + net = PSRN(2, ['Add', 'Identity'], 1, device='cpu') + x = torch.linspace(-1, 1, 32, dtype=torch.float32).reshape(16, 2) + for shape in ((16,), (16, 1)): + y = torch.linspace(-.5, .5, 16, dtype=torch.float64).reshape(shape) + bound = net.estimate_dense_score_bytes(x, y) + for budget in (bound - 1, bound): + modern = net.score_mse(x, y, workspace_budget_bytes=budget) + modern_route = net.last_score_strategy + with patch.object(torch.Tensor, 'dtype', new=property(legacy_tensor_dtype)), \ + patch.object(torch, 'result_type', side_effect=lambda *a: LegacyDtype(real_result_type(*a))), \ + patch.object(torch, 'get_default_dtype', side_effect=lambda: LegacyDtype(real_default_dtype())), \ + patch.object(torch, 'empty', side_effect=legacy_empty): + self.assertEqual(net.estimate_dense_score_bytes(x, y), bound) + legacy = net.score_mse(x, y, workspace_budget_bytes=budget) + self.assertEqual(net.last_score_strategy, modern_route) + self.assertTrue(torch.equal(integer_view(legacy), integer_view(modern))) + self.assertEqual(len(allocations), 24) + self.assertIn(torch.float64, allocations) + self.assertIn(torch.float32, allocations) + + def test_active_autocast_overrides_auto_and_explicit_tiles(self): + net = PSRN(2, ['Identity', 'Pow2'], 1, device='cpu') + x, y = torch.tensor([[1., .5], [.25, 2.]]), torch.tensor([2 ** -12, -.5]) + expected = original_scores(net, x, y) + for strategy in ('auto', 'dense', 'tiled'): + with patch.object(torch, 'is_autocast_enabled', return_value=True) as query, \ + patch.object(net, 'estimate_dense_score_bytes', side_effect=AssertionError('AMP estimate')), \ + patch.object(net.list[-1], 'iter_tiles', side_effect=AssertionError('AMP tiles')): + actual = net.score_mse(x, y, strategy=strategy, workspace_budget_bytes=1) + query.assert_called_once_with('cpu') + self.assertEqual(net.last_score_strategy, 'dense') + self.assertTrue(torch.equal(integer_view(actual), integer_view(expected))) + + def test_actual_cpu_autocast_preserves_cat_promotion_and_scores(self): + net = PSRN(2, ['Identity', 'Pow2', 'Add', 'Sin'], 2, device='cpu') + for dtype in (torch.bfloat16, torch.float32, torch.float64): + x = torch.tensor([[1., .5], [.25, 2.], [-1., 0.]], dtype=dtype) + for shape in ((3,), (3, 1)): + y = torch.tensor([2 ** -12, -.5, .25], dtype=dtype).reshape(shape) + with torch.autocast('cpu', dtype=torch.bfloat16): + expected = original_scores(net, x, y) + for strategy in ('auto', 'dense', 'tiled'): + actual = net.score_mse(x, y, strategy=strategy, workspace_budget_bytes=1) + self.assertEqual(net.last_score_strategy, 'dense') + self.assertTrue(torch.equal(integer_view(actual), integer_view(expected))) + + def test_off_autocast_queries_once_and_keeps_original_routes(self): + net = PSRN(2, ['Identity', 'Add'], 1, device='cpu') + x, y = torch.ones(17, 2), torch.zeros(17) + expected = original_scores(net, x, y) + for strategy, route in (('dense', 'dense'), ('tiled', 'tiled'), ('auto', 'tiled')): + with patch.object(torch, 'is_autocast_enabled', wraps=torch.is_autocast_enabled) as query: + actual = net.score_mse(x, y, strategy=strategy, workspace_budget_bytes=1) + query.assert_called_once_with('cpu') + self.assertEqual(net.last_score_strategy, route) + self.assertTrue(torch.equal(integer_view(actual), integer_view(expected))) + + def test_half_operator_dtype_counterexample_without_cuda(self): + # Emulate the mixed output dtypes of CUDA AMP Identity and Pow2. + # This is dtype arithmetic on CPU, not a vendor CUDA AMP runtime test. + identity = torch.tensor([[1.]], dtype=torch.float16) + power = identity.float().pow(2) + target = torch.tensor([2 ** -12], dtype=torch.float16) + dense = (torch.cat((identity, power), dim=1) - target).pow(2).float() + per_operator = torch.cat(((identity - target).pow(2).float(), + (power - target).pow(2).float()), dim=1) + self.assertEqual(dense[0, 0].item(), 0.9995117783546448) + self.assertEqual(per_operator[0, 0].item(), 1.) + self.assertFalse(torch.equal(integer_view(dense), integer_view(per_operator))) + + def test_integer_cat_promotion_counterexample_and_compatibility_routes(self): + net = PSRN(1, ['Identity', 'Sin'], 1, device='cpu') + x = torch.tensor([[2 ** 24 + 1]], dtype=torch.int64) + y = torch.zeros(1, dtype=torch.int64) + expected = original_scores(net, x, y) + per_operator_square = x.pow(2).float().reshape(-1) + self.assertEqual(expected[0].item(), float(2 ** 48)) + self.assertEqual(per_operator_square[0].item(), float(2 ** 48 + 2 ** 25)) + self.assertFalse(torch.equal(integer_view(expected[:1]), integer_view(per_operator_square))) + for strategy in ('auto', 'dense', 'tiled'): + actual = net.score_mse(x, y, strategy=strategy, workspace_budget_bytes=1) + self.assertEqual(net.last_score_strategy, 'dense') + self.assertTrue(torch.equal(integer_view(actual), integer_view(expected))) + for dtype, operators, samples, target_dtype, rank in itertools.product( + (torch.int32, torch.int64), + (['Identity', 'Sin'], ['Add', 'Identity', 'Div', 'Log']), + (1, 3, 17), (torch.int64, torch.float32, torch.float64), (1, 2)): + net = PSRN(2, operators, 1, device='cpu') + storage = torch.arange(samples * 8, dtype=dtype).reshape(samples * 2, 4) % 5 + 1 + x = storage[::2, ::2] + y = torch.arange(samples, dtype=target_dtype) % 3 + if rank == 2: + y = y.reshape(samples, 1) + expected = original_scores(net, x, y) + for strategy in ('auto', 'dense', 'tiled'): + with patch.object(net, 'estimate_dense_score_bytes', side_effect=AssertionError('integer estimate')), \ + patch.object(net.list[-1], 'iter_tiles', side_effect=AssertionError('integer tiles')): + actual = net.score_mse(x, y, strategy=strategy, workspace_budget_bytes=1) + self.assertEqual(net.last_score_strategy, 'dense') + self.assertTrue(torch.equal(integer_view(actual), integer_view(expected))) + + def test_complex_inputs_preserve_original_scoring_error(self): + net = PSRN(2, ['Identity', 'Abs'], 1, device='cpu') + x = torch.tensor([[1 + 2j, -1j]]) + y = torch.zeros(1) + with self.assertRaises(RuntimeError) as reference: + original_scores(net, x, y) + for strategy in ('auto', 'dense', 'tiled'): + with self.assertRaises(RuntimeError) as actual: + net.score_mse(x, y, strategy=strategy, workspace_budget_bytes=1) + self.assertEqual(str(actual.exception), str(reference.exception)) + self.assertEqual(net.last_score_strategy, 'dense') + + def test_off_amp_all_operator_float_output_dtypes(self): + # Check each output separately: cat would hide heterogeneous outputs. + for dtype, operator in itertools.product( + (torch.float16, torch.bfloat16, torch.float32, torch.float64), FAMILIES[-1]): + net = PSRN(2, [operator], 1, device='cpu') + x = torch.tensor([[.5, 1.25]], dtype=dtype) + with torch.autocast('cpu', enabled=False): + self.assertEqual(net.forward(x).dtype, dtype, (dtype, operator)) + + def test_no_grad_default_accumulator_dtype_and_three_layer_mask(self): + operators = FAMILIES[-1] + prefix = PSRN(2, operators, 2, device='cpu') + mask = torch.arange(prefix.out_dim) < 3 + net = PSRN(2, operators, 3, dr_mask=mask, device='cpu') + x = torch.tensor([[.1, .3], [.6, 1.2]], requires_grad=True) + y = torch.tensor([.5, -.25], dtype=torch.float64, requires_grad=True) + previous = torch.get_default_dtype() + try: + for accumulator_dtype in (torch.float32, torch.float64): + torch.set_default_dtype(accumulator_dtype) + actual = net.score_mse(x, y, strategy='dense') + expected = original_scores(net, x, y) + self.assertEqual(actual.dtype, accumulator_dtype) + self.assertFalse(actual.requires_grad) + self.assertTrue(torch.equal(integer_view(actual), integer_view(expected))) + finally: + torch.set_default_dtype(previous) + + +if __name__ == '__main__': + unittest.main() diff --git a/test/test_row_aligned_tiles.py b/test/test_row_aligned_tiles.py new file mode 100644 index 0000000..0878a63 --- /dev/null +++ b/test/test_row_aligned_tiles.py @@ -0,0 +1,190 @@ +"""Exact row-tile geometry and scoring checks; no accelerator performance claim.""" +import importlib +from pathlib import Path +import sys +import types +import unittest +from unittest.mock import patch + +import torch + +package = types.ModuleType('_psrn_rowtile_tests') +package.__path__ = [str(Path(__file__).resolve().parents[1] / 'psrn')] +sys.modules[package.__name__] = package +models = importlib.import_module(package.__name__ + '.model.models') +PSRN, SymbolLayer, DRLayer = models.PSRN, models.SymbolLayer, models.DRLayer +OPS = ['Add', 'Mul', 'Sub', 'Div', 'SemiSub', 'SemiDiv', 'Pow', 'Identity', + 'Sin', 'Cos', 'Exp', 'Log', 'Neg', 'Inv', 'Sign', 'Pow2', 'Pow3', + 'Sigmoid', 'Abs', 'Cosh', 'Tanh', 'Sqrt'] + + +def original_scores(net, x, y): + with torch.no_grad(): + sums = torch.zeros((1, net.out_dim), device=net.device) + for row in range(x.shape[0]): + h = net.forward(x[row].reshape(1, -1)) + target = y[row] + if target.dim() > 1: + target = target.reshape(1) + diff = h - target + sums += diff ** 2 + scores = (sums / x.shape[0]).reshape(-1) + scores[~torch.isfinite(scores)] = float('inf') + return scores + + +def dense_expressions(net, indices): + """Recover expressions through the original dense operand tables.""" + tables = {id(layer): layer.get_offset_tensor('cpu') for layer in net.list + if isinstance(layer, SymbolLayer)} + + def expression(index, position): + if position < 0: + return net.current_expr_ls[index] + layer = net.list[position] + if isinstance(layer, DRLayer): + return expression(layer.dr_indices[index].item(), position - 1) + start = 0 + for op in layer.list: + if index < start + op.out_dim: + break + start += op.out_dim + left, right = tables[id(layer)][index].tolist() + if op.is_unary: + return op.operator.get_expr(expression(left, position - 1)) + return op.operator.get_expr(expression(left, position - 1), + expression(right, position - 1)) + + return [expression(index, len(net.list) - 1) for index in indices] + + +class RowAlignedTileTests(unittest.TestCase): + @classmethod + def setUpClass(cls): + torch.set_num_threads(1) + torch.manual_seed(20261005) + + def assert_prediction_bits(self, actual, expected): + self.assertTrue(torch.equal(torch.isnan(actual), torch.isnan(expected))) + self.assertTrue(torch.equal(torch.isposinf(actual), torch.isposinf(expected))) + self.assertTrue(torch.equal(torch.isneginf(actual), torch.isneginf(expected))) + finite = torch.isfinite(expected) + integer_dtype = torch.int32 if actual.dtype == torch.float32 else torch.int64 + self.assertTrue(torch.equal(actual.contiguous().view(integer_dtype)[finite], + expected.contiguous().view(integer_dtype)[finite])) + + def collect(self, layer, x, cap): + values, next_offset = [], 0 + for offset, tile in layer.iter_tiles(x, cap): + self.assertEqual(offset, next_offset) + self.assertLessEqual(tile.shape[1], cap) + self.assertEqual(tile.shape[0], x.shape[0]) + values.append(tile) + next_offset += tile.shape[1] + self.assertEqual(next_offset, layer.out_dim) + return torch.cat(values, dim=1) + + def test_exhaustive_triangle_row_boundaries_and_fallback(self): + for width in range(1, 21): + layer = SymbolLayer(width, ['Add', 'SemiSub'], device='cpu') + x = torch.arange(2 * width, dtype=torch.float64).reshape(2, width) + expected = torch.cat([layer(row.reshape(1, -1)) for row in x]) + for cap in {1, 7, 31, max(1, width - 1), width, width + 1, + 2 * width + 1, 1048576}: + with self.subTest(width=width, cap=cap): + with patch.object(torch, 'searchsorted', wraps=torch.searchsorted) as search: + actual = self.collect(layer, x, cap) + self.assert_prediction_bits(actual, expected) + self.assertEqual(search.call_count == 0, width <= cap) + + def test_all_operators_raw_predictions_reordered_and_strided(self): + orders = [OPS, list(reversed(OPS)), OPS[::2] + OPS[1::2]] + for dtype in (torch.float32, torch.float64): + for width in (1, 3, 8): + storage = torch.rand(6, width * 2, dtype=dtype) * 4 - 2 + storage[0, 0] = 0 + x = storage[::2, ::2] + for operators in orders: + layer = SymbolLayer(width, operators, device='cpu') + expected = torch.cat([layer(row.reshape(1, -1)) for row in x]) + for cap in {1, 7, 31, width, width + 1, 1048576}: + with self.subTest(dtype=dtype, width=width, operators=operators, cap=cap): + # CPU Pow historically differs for some gathered + # layouts: production scoring retains its dense + # override. Its predictions are checked separately. + actual = self.collect(layer, x, cap) + non_pow = torch.ones(layer.out_dim, dtype=torch.bool) + start = 0 + for op in layer.list: + if isinstance(op, models.Pow): + non_pow[start:start + op.out_dim] = False + start += op.out_dim + self.assert_prediction_bits(actual[:, non_pow], expected[:, non_pow]) + + def check_scores(self, net, x, y, caps): + net.current_expr_ls = ['x{}'.format(i) for i in range(net.n_variables)] + expected = original_scores(net, x, y) + k = min(5, net.out_dim) + selected = torch.topk(expected, k, largest=False).indices + exprs = dense_expressions(net, selected.tolist()) + for cap in caps: + actual = net.score_mse(x, y, strategy='tiled', tile_size=cap, + sample_chunk_size=2) + self.assertTrue(torch.equal(actual.view(torch.int32), expected.view(torch.int32))) + self.assertTrue(torch.equal(torch.topk(actual, k, largest=False).indices, selected)) + self.assertEqual([net.get_expr(i) for i in selected.tolist()], exprs) + has_pow = any(isinstance(op, models.Pow) for op in net.list[-1].list) + self.assertEqual(net.last_score_strategy, 'dense' if has_pow else 'tiled') + + def test_each_operator_two_three_layers_mask_dtypes_target_shapes(self): + for operator in OPS: + for layers in (2, 3): + prefix = PSRN(2, [operator], layers - 1, device='cpu') + mask = torch.arange(prefix.out_dim) % 3 != 1 + for masked in (False, True): + net = PSRN(2, [operator], layers, dr_mask=mask if masked else None, + device='cpu') + for dtype in (torch.float32, torch.float64): + x = torch.tensor([[0., 9., -1., 9.], [2., 9., 0.5, 9.], + [-0.5, 9., 1.2, 9.]], dtype=dtype)[:, ::2] + for target_shape in ((3,), (3, 1)): + y = torch.tensor([1., 0., -1.], dtype=torch.float64).reshape(target_shape) + with self.subTest(op=operator, layers=layers, mask=masked, + dtype=dtype, target_shape=target_shape): + self.check_scores(net, x, y, (7, 31, 1048576)) + + def test_combined_reordered_masked_networks_and_nonfinite_inputs(self): + non_pow = [op for op in OPS if op != 'Pow'] + for operators in (OPS, list(reversed(OPS)), OPS[::2] + OPS[1::2], + non_pow, list(reversed(non_pow)), non_pow[::2] + non_pow[1::2]): + for layers in (2, 3): + prefix = PSRN(2, operators, layers - 1, device='cpu') + # Bound the combined three-layer final width without changing + # the preceding full layer's operator evaluation or ordering. + mask = torch.arange(prefix.out_dim) < 5 + for dtype in (torch.float32, torch.float64): + net = PSRN(2, operators, layers, dr_mask=mask, device='cpu') + x = torch.tensor([[0., 9., -1., 9.], [float('inf'), 9., 0.5, 9.], + [float('nan'), 9., -float('inf'), 9.]], dtype=dtype)[:, ::2] + y = torch.tensor([1., 0., -1.], dtype=dtype).reshape(3, 1, 1, 1) + self.check_scores(net, x, y, (1, 7, 31, 1048576)) + finite_x = torch.tensor([[0.5, 1.2], [0.8, 1.1], [1.3, 0.7]], dtype=dtype) + self.check_scores(net, finite_x, y.reshape(3), (7, 31, 1048576)) + + def test_row_tiles_cross_sample_chunk_tail_without_pow_override(self): + for operators in (['Add', 'Sub', 'SemiDiv', 'Identity'], + ['Div', 'SemiSub', 'Mul', 'Cos']): + net = PSRN(3, operators, 2, device='cpu') + x = torch.rand(34, 6, dtype=torch.float64)[::2, ::2] * 2 - 1 + y = torch.linspace(-1, 1, 17).reshape(17, 1) + expected = original_scores(net, x, y) + width = net.list[-1].in_dim + for cap in (width - 1, width, width + 1, width * 2 + 1, 1048576): + actual = net.score_mse(x, y, strategy='tiled', tile_size=cap, + sample_chunk_size=16) + self.assertEqual(net.last_score_strategy, 'tiled') + self.assertTrue(torch.equal(actual.view(torch.int32), expected.view(torch.int32))) + + +if __name__ == '__main__': + unittest.main() diff --git a/test/test_scoring_dispatch.py b/test/test_scoring_dispatch.py new file mode 100644 index 0000000..6455ce7 --- /dev/null +++ b/test/test_scoring_dispatch.py @@ -0,0 +1,166 @@ +"""Run with python -m unittest discover -s test -p test_scoring_dispatch.py. + +Loads only model code so the scoring tests do not require regressor extras. +""" +import importlib +from pathlib import Path +import sys +import types +import unittest +from unittest.mock import patch + +import torch + +PACKAGE = "_psrn_scoring_tests" +package = types.ModuleType(PACKAGE) +package.__path__ = [str(Path(__file__).resolve().parents[1] / "psrn")] +sys.modules.setdefault(PACKAGE, package) +PSRN = importlib.import_module(PACKAGE + ".model.models").PSRN + + +def original_scores(net, x, y): + with torch.no_grad(): + sums = torch.zeros((1, net.out_dim), device=net.device) + for row in range(x.shape[0]): + prediction = net.forward(x[row].reshape(1, -1)) + target = y[row] + if target.dim() > 1: + target = target.reshape(1) + diff = prediction - target + sums += diff ** 2 + scores = (sums / x.shape[0]).reshape(-1) + scores[~torch.isfinite(scores)] = float("inf") + return scores + + +class ScoringDispatchTests(unittest.TestCase): + def test_budget_boundary_routes_and_exact_scores(self): + net = PSRN(2, ["Add", "Sub", "Identity"], 1, device="cpu") + for samples, chunk in ((1, 16), (16, 16), (17, 16), (100, 16), + (3, 2), (1, 1), (17, 32)): + with self.subTest(samples=samples, chunk=chunk): + x = torch.arange(samples * 2, dtype=torch.float32).reshape(samples, 2) / 7 + y = torch.linspace(-1, 1, samples).reshape(samples, 1) + expected = original_scores(net, x, y) + bound = net.estimate_dense_score_bytes(x, y) + groups = (samples + min(chunk, samples) - 1) // min(chunk, samples) + # Smallest integer base budget admitting the dense estimate. + scale = chunk * groups + budget = (bound * samples + scale - 1) // scale + with patch.object(net.list[-1], "iter_tiles", wraps=net.list[-1].iter_tiles) as tiles: + actual = net.score_mse(x, y, sample_chunk_size=chunk, + workspace_budget_bytes=budget) + self.assertEqual(net.last_score_strategy, "dense") + tiles.assert_not_called() + self.assertTrue(torch.equal(actual, expected)) + actual = net.score_mse(x, y, tile_size=3, sample_chunk_size=chunk, + workspace_budget_bytes=budget - 1) + self.assertEqual(net.last_score_strategy, "tiled") + self.assertGreater(tiles.call_count, 0) + self.assertTrue(torch.equal(actual, expected)) + + def test_small_queries_use_soft_allowance_above_base_budget(self): + net = PSRN(1, ["Identity"], 1, device="cpu") + base = 512 * 1024 ** 2 + # Exercise a just-over-base estimate without allocating a huge model. + with patch.object(net, "estimate_dense_score_bytes", return_value=base + 1): + for samples, expected_strategy in ((1, "dense"), (16, "tiled"), + (17, "dense"), (100, "dense")): + x = torch.ones(samples, 1) + y = torch.zeros(samples) + actual = net.score_mse(x, y) + self.assertEqual(net.last_score_strategy, expected_strategy) + self.assertTrue(torch.equal(actual, original_scores(net, x, y))) + + def test_scalar_promotion_changes_estimate_and_routing(self): + net = PSRN(2, ["Identity"], 1, device="cpu") + x = torch.tensor([[.5, 1.25], [-.75, 2.5]], dtype=torch.float32).repeat(8, 1) + scalar_targets = torch.linspace(-1, 1, 32, dtype=torch.float64)[::2] + for shape, expected_bound, expected_strategy in ( + ((16,), 32, "dense"), ((16, 1), 48, "tiled"), + ((16, 1, 1, 1), 48, "tiled")): + y = scalar_targets.reshape(shape) + with self.subTest(shape=shape): + self.assertEqual(net.estimate_dense_score_bytes(x, y), expected_bound) + expected = original_scores(net, x, y) + actual = net.score_mse(x, y, workspace_budget_bytes=32) + self.assertEqual(net.last_score_strategy, expected_strategy) + self.assertTrue(torch.equal(actual, expected)) + for strategy in ("dense", "tiled"): + actual = net.score_mse(x, y, strategy=strategy, tile_size=1) + self.assertTrue(torch.equal(actual, expected)) + + def test_estimator_observes_module_replacement_and_default_dtype(self): + net = PSRN(1, ["Identity"], 1, device="cpu") + x, y = torch.ones(16, 1, dtype=torch.float32), torch.zeros(16, dtype=torch.float32) + original_bound = net.estimate_dense_score_bytes(x, y) + self.assertEqual(original_bound, 16) + net.to(dtype=torch.float64) + self.assertEqual(net.estimate_dense_score_bytes(x.double(), y.double()), 28) + net.list[-1] = PSRN(1, ["Add"], 1, device="cpu").list[-1] + self.assertEqual(net.estimate_dense_score_bytes(x, y), 28) + actual = net.score_mse(x, y, workspace_budget_bytes=original_bound) + self.assertEqual(net.last_score_strategy, "tiled") + self.assertTrue(torch.equal(actual, original_scores(net, x, y))) + previous_dtype = torch.get_default_dtype() + try: + torch.set_default_dtype(torch.float64) + self.assertEqual(net.estimate_dense_score_bytes(x, y), 32) + actual = net.score_mse(x, y, strategy="dense") + self.assertEqual(actual.dtype, torch.float64) + self.assertTrue(torch.equal(actual, original_scores(net, x, y))) + finally: + torch.set_default_dtype(previous_dtype) + + def test_forced_tiles_mask_noncontiguous_fp64_and_target_shapes(self): + operators = ["Add", "SemiSub", "SemiDiv", "Identity", "Sin", "Inv", "Log"] + pre = PSRN(2, operators, 1, device="cpu") + mask = torch.arange(pre.out_dim) % 3 != 0 + net = PSRN(2, operators, 2, dr_mask=mask, device="cpu") + x = torch.tensor([[0., 9., -1., 9.], [2., 9., 3., 9.], + [-2., 9., 0.5, 9.]], dtype=torch.float64)[:, ::2] + for target_dtype in (torch.float32, torch.float64): + for shape in ((3,), (3, 1), (3, 1, 1), (3, 1, 1, 1)): + y = torch.tensor([1., 0., -1.], dtype=target_dtype).reshape(shape) + expected = original_scores(net, x, y) + for strategy in ("dense", "tiled", "auto"): + actual = net.score_mse(x, y, tile_size=11, sample_chunk_size=2, + strategy=strategy) + self.assertTrue(torch.equal(actual, expected), (strategy, shape, target_dtype)) + self.assertTrue(torch.equal(torch.topk(actual, 5, largest=False).indices, + torch.topk(expected, 5, largest=False).indices)) + self.assertTrue(all(layer.offset_tensor is None for layer in net.list + if hasattr(layer, "offset_tensor"))) + + def test_cpu_pow_remains_dense_even_when_tiles_requested(self): + net = PSRN(3, ["Pow", "Identity"], 1, device="cpu") + x = torch.tensor([[0., -2., 0.5], [1.1, 2.3, -0.7]]) + y = torch.zeros(2) + with patch.object(net.list[-1], "iter_tiles", side_effect=AssertionError("CPU Pow tiled")): + actual = net.score_mse(x, y, strategy="tiled", workspace_budget_bytes=1) + self.assertEqual(net.last_score_strategy, "dense") + self.assertTrue(torch.equal(actual, original_scores(net, x, y))) + + def test_bound_tracks_dtype_and_triangle_storage(self): + net = PSRN(10, ["Add"], 1, device="cpu") + x = torch.ones(1, 10) + y = torch.zeros(1) + # FP32 scores + input + results + two int32 indices + two gathers. + expected = 55 * 4 + 10 * 4 + 55 * (4 + 8 + 2 * 4) + self.assertEqual(net.estimate_dense_score_bytes(x, y), expected) + self.assertGreater(net.estimate_dense_score_bytes(x.double(), y.double()), expected) + + def test_invalid_dispatch_options_and_scalar_target(self): + net = PSRN(2, ["Add"], 1, device="cpu") + x, y = torch.ones(1, 2), torch.zeros(1) + for options in ({"strategy": "unknown"}, {"workspace_budget_bytes": True}, + {"workspace_budget_bytes": 0}, {"tile_size": 0}, + {"sample_chunk_size": False}): + with self.assertRaises(ValueError): + net.score_mse(x, y, **options) + with self.assertRaises(ValueError): + net.score_mse(x, y[0]) + + +if __name__ == "__main__": + unittest.main()