diff --git a/helion/_compiler/cute/aux_tensor.py b/helion/_compiler/cute/aux_tensor.py index 8964a890b..7ae66b4a1 100644 --- a/helion/_compiler/cute/aux_tensor.py +++ b/helion/_compiler/cute/aux_tensor.py @@ -94,9 +94,9 @@ class Tcgen05AuxTensorDescriptor: the host tensor's FX node — the concrete shape / dtype / device the producer-body codegen needs at MMA-codegen time (the TMA atom and SMEM ring are sized from these). - - ``broadcast_axis``: ``None`` for the exact-shape rank-2 form - (``residual[tile_m, tile_n]``); ``1`` for the trailing-axis - rowvec broadcast form (``bias[tile_n]``). Matches the field on + - ``broadcast_axis``: ``None`` for exact-shape loads, ``1`` for a + trailing-axis row vector, and ``2`` for an explicit rank-1 column + vector. Matches the field on :class:`_AuxiliaryTensorStep` so the productive body uses the same axis convention as the existing per-thread splice. - ``store_value_node``: the FX node of the store value whose diff --git a/helion/_compiler/cute/cute_epilogue.py b/helion/_compiler/cute/cute_epilogue.py index 1608c9f39..949d1e32d 100644 --- a/helion/_compiler/cute/cute_epilogue.py +++ b/helion/_compiler/cute/cute_epilogue.py @@ -22,11 +22,11 @@ - Auxiliary-tensor binary ops: same shape as above, but one or more steps are ``add/sub/mul/div`` with the chain carrier as one operand and a ``helion.language.load(aux_tensor, [...])`` - call as the other. Two aux load shapes are accepted: the - exact-shape rank-2 form (``residual[tile_m, tile_n]``) and the - rank-1 trailing-axis (rowvec) broadcast form (``bias[tile_n]``). + call as the other. Accepted shapes include exact-shape rank-2 loads, + rank-1 trailing-axis row vectors, and explicit rank-1 column vectors + spelled ``bias[tile_m].unsqueeze(-1)``. See :class:`_AuxiliaryTensorStep` for the canonical contract. - Forms outside these two — 3-D underlying tensors with a static + Other forms — 3-D underlying tensors with a static collapse, mismatched indices, leading-axis rank-1 (``bias[tile_m]``), kwargs — are rejected to the loud-failure backstop. @@ -120,7 +120,8 @@ class _AuxiliaryTensorStep: ``broadcast_axis`` is ``None`` when the aux tensor matches the carrier rank exactly (``residual[tile_m, tile_n]``), ``1`` for a trailing-axis (rowvec) broadcast aux load (``bias[tile_n]`` - with shape ``(N,)``), or ``0`` for a leading-axis (row / + with shape ``(N,)``), ``2`` for an explicit column vector + (``bias[tile_m].unsqueeze(-1)``), or ``0`` for a leading-axis (row / M-axis) broadcast aux load spelled explicitly as a ``(1, N)`` tensor indexed ``bias[tile_m, tile_n]``. A *bare* rank-1 operand on the RHS (``acc + bias[tile_m]``) is still rejected @@ -132,11 +133,10 @@ class _AuxiliaryTensorStep: M axis themselves, so row 0 broadcasts over every output row. The splice site (``memory_ops._codegen_cute_store_tcgen05_tile``) owns the - canonical broadcast-view contract — for both broadcast axes it - builds a 2-D logical view with stride 0 on the broadcast - (M) axis and stride 1 on the data (N) axis so the existing - ``partition_C → flat_divide → partition_D`` pipeline can run - unchanged. Mirrors Quack's ``RowVecLoad`` epilogue + canonical broadcast-view contract. Row vectors use stride ``(0, 1)``; + compact column vectors use stride ``(1, 0)``. Both feed the existing + ``partition_C → flat_divide → partition_D`` pipeline unchanged. The + row-vector form mirrors Quack's ``RowVecLoad`` epilogue (``quack/quack/epi_ops.py``). """ @@ -370,65 +370,80 @@ def _is_helion_load_node(node: torch.fx.Node) -> bool: return node.op == "call_function" and node.target is helion_load -def _canonical_aux_load_operand(node: torch.fx.Node) -> tuple[torch.fx.Node, str]: - """Return the underlying aux load plus its local-value template. +@dataclasses.dataclass(frozen=True) +class _AuxLoadOperand: + load_node: torch.fx.Node + value_template: str = "{aux}" + rank1_data_axis: int | None = None + + +def _aux_load_operand(node: torch.fx.Node) -> _AuxLoadOperand | None: + """Normalize the accepted wrappers around one auxiliary load. - Canonical unwrapping currently accepts raw aux loads and fp32 casts of - aux loads; other dtype casts stay outside the fused aux pattern. + ``rank1_data_axis`` records the surviving axis after an explicit + rank-1 ``unsqueeze``. This distinguishes a real column-vector expression + from a bare rank-1 operand, whose broadcast direction follows PyTorch's + trailing-axis rule. """ if _is_helion_load_node(node): - return node, "{aux}" + return _AuxLoadOperand(node) + if node.op != "call_function" or node.kwargs: + return None + if ( - node.op == "call_function" - and node.target is torch.ops.prims.convert_element_type.default - and not node.kwargs + node.target is torch.ops.prims.convert_element_type.default and len(node.args) == 2 and node.args[1] is torch.float32 and isinstance(node.args[0], torch.fx.Node) ): - inner = node.args[0] - if _is_helion_load_node(inner): - return inner, "({aux}).to(cutlass.Float32)" - return node, "{aux}" - + inner = _aux_load_operand(node.args[0]) + if inner is None: + return None + return dataclasses.replace( + inner, value_template=f"({inner.value_template}).to(cutlass.Float32)" + ) -def _aux_load_operand(node: torch.fx.Node) -> tuple[torch.fx.Node, str]: - """Return ``(load_node, aux_template)`` for accepted aux operands. + if ( + node.target is torch.ops.aten.unsqueeze.default + and len(node.args) == 2 + and isinstance(node.args[0], torch.fx.Node) + and isinstance(node.args[1], int) + and not isinstance(node.args[1], bool) + ): + inner = _aux_load_operand(node.args[0]) + load_val = inner.load_node.meta.get("val") if inner is not None else None + dim = node.args[1] % 2 + if ( + inner is None + or inner.rank1_data_axis is not None + or not isinstance(load_val, torch.Tensor) + or load_val.ndim != 1 + or node.args[1] not in (-2, -1, 0, 1) + ): + return None + return dataclasses.replace(inner, rank1_data_axis=1 - dim) - ``aux_template`` is a ``{aux}``-keyed expression applied to the loaded - per-thread aux local before the outer binary op. This covers wrapped - residual terms such as ``0.5 * residual[tile_m, tile_n].to(torch.float32)`` - without broadening the general chain model to arbitrary non-scalar binary - reuse. - """ - load_node, aux_template = _canonical_aux_load_operand(node) - if load_node is not node: - return load_node, aux_template - if node.op != "call_function" or node.target is not torch.ops.aten.mul.Tensor: - return node, "{aux}" - if node.kwargs or len(node.args) < 2: - return node, "{aux}" - lhs = node.args[0] - rhs = node.args[1] + if node.target is not torch.ops.aten.mul.Tensor or len(node.args) < 2: + return None + lhs, rhs = node.args[:2] if isinstance(lhs, torch.fx.Node): scalar = _extract_scalar(rhs) - inner = lhs + inner_node = lhs elif isinstance(rhs, torch.fx.Node): scalar = _extract_scalar(lhs) - inner = rhs + inner_node = rhs else: - return node, "{aux}" - if scalar is None: - return node, "{aux}" - load_node, inner_template = _canonical_aux_load_operand(inner) - if load_node is inner and not _is_helion_load_node(load_node): - return node, "{aux}" - return load_node, f"(({inner_template}) * {scalar!r})" + return None + inner = _aux_load_operand(inner_node) + if scalar is None or inner is None: + return None + return dataclasses.replace( + inner, value_template=f"(({inner.value_template}) * {scalar!r})" + ) def _is_aux_load_operand_node(node: torch.fx.Node) -> bool: - load_node, _ = _aux_load_operand(node) - return _is_helion_load_node(load_node) + return _aux_load_operand(node) is not None @dataclasses.dataclass(frozen=True) @@ -749,32 +764,36 @@ def _classify_binary( # tensor load; the other is the chain carrier. If both look # like aux loads or neither does, bail — the chain has no # unique carrier. - lhs_load, lhs_aux_template = _aux_load_operand(lhs) - rhs_load, rhs_aux_template = _aux_load_operand(rhs) - lhs_kind = aux_tensor_load_kind( - lhs_load, - carrier_tile_shape=carrier_tile_shape, - carrier_tile_index_nodes=carrier_tile_index_nodes, - carrier_global_shape=carrier_global_shape, - ) - rhs_kind = aux_tensor_load_kind( - rhs_load, - carrier_tile_shape=carrier_tile_shape, - carrier_tile_index_nodes=carrier_tile_index_nodes, - carrier_global_shape=carrier_global_shape, - ) + lhs_operand = _aux_load_operand(lhs) + rhs_operand = _aux_load_operand(rhs) + + def load_kind(operand: _AuxLoadOperand | None) -> tuple[str, int | None] | None: + if operand is None: + return None + return aux_tensor_load_kind( + operand.load_node, + carrier_tile_shape=carrier_tile_shape, + carrier_tile_index_nodes=carrier_tile_index_nodes, + carrier_global_shape=carrier_global_shape, + explicit_rank1_data_axis=operand.rank1_data_axis, + ) + + lhs_kind = load_kind(lhs_operand) + rhs_kind = load_kind(rhs_operand) aux_load: torch.fx.Node carrier: torch.fx.Node forward_form: bool if lhs_kind is not None and rhs_kind is None: - aux_load = lhs_load - aux_template = lhs_aux_template + assert lhs_operand is not None + aux_load = lhs_operand.load_node + aux_template = lhs_operand.value_template carrier = rhs forward_form = False # carrier is the right operand aux_kind = lhs_kind elif rhs_kind is not None and lhs_kind is None: - aux_load = rhs_load - aux_template = rhs_aux_template + assert rhs_operand is not None + aux_load = rhs_operand.load_node + aux_template = rhs_operand.value_template carrier = lhs forward_form = True # carrier is the left operand aux_kind = rhs_kind @@ -968,9 +987,8 @@ def analyze_tcgen05_unary_epilogue_chain( (``add`` / ``sub`` / ``mul`` / ``div`` against a compile-time Python literal), and auxiliary-tensor binary in two forms: exact-shape (``residual[tile_m, tile_n]``, rank-2 aux matching - the carrier tile shape) and rank-1 trailing-axis (rowvec) - broadcast (``bias[tile_n]``, where the single load index - symbol matches the carrier's trailing tile-id symbol). Other + the carrier tile shape), rank-1 trailing-axis rowvec broadcast, + and explicit rank-1 column-vector broadcast via ``unsqueeze``. Other shapes — 3-D collapsed loads, indices that are not exactly the carrier trailing tile-id symbol, leading-axis rank-1 (``bias[tile_m]``), kwargs — are rejected so the loud-failure diff --git a/helion/_compiler/cute/cute_fx_walk.py b/helion/_compiler/cute/cute_fx_walk.py index eff0c8e70..aac55dabe 100644 --- a/helion/_compiler/cute/cute_fx_walk.py +++ b/helion/_compiler/cute/cute_fx_walk.py @@ -195,6 +195,7 @@ def aux_tensor_load_kind( carrier_tile_shape: tuple[object, ...] | None, carrier_tile_index_nodes: tuple[torch.fx.Node, ...] | None = None, carrier_global_shape: tuple[object, ...] | None = None, + explicit_rank1_data_axis: int | None = None, ) -> tuple[str, int | None] | None: """Classify ``node`` as a recognized auxiliary-tensor load kind, or return ``None`` when the load is not eligible for fusion. @@ -217,9 +218,9 @@ def aux_tensor_load_kind( ``acc + bias[tile_m]`` is either a shape error (BM ≠ BN) or rowvec broadcast against the trailing axis (BM == BN). A user wanting M-axis (column-vector) broadcasting via a bare - rank-1 operand must write it explicitly (``bias[tile_m][:, - None]`` / ``.unsqueeze(-1)``); that is a separate, deferred - pattern. See :class:`Tcgen05UnaryEpilogueChain` + rank-1 operand must write it explicitly with ``.unsqueeze(-1)``; + the operand normalizer records that unambiguous direction. See + :class:`Tcgen05UnaryEpilogueChain` (``cute_epilogue.py``) for the splice-side broadcast-view contract. - ``("broadcast", 0)``: a rank-2 leading-axis broadcast aux @@ -249,6 +250,11 @@ def aux_tensor_load_kind( residual has a non-zero trailing stride and falls through to the exact matcher. + An explicit rank-1 ``unsqueeze`` is represented by + ``explicit_rank1_data_axis``. Axis 0 is a column vector and reuses the + ``("broadcast", 2)`` codegen; axis 1 is the explicit spelling of the + ordinary row-vector case. + A rank-3 exact residual is accepted when its leading index is the carrier's leading passthrough tile. The classifier returns ``None`` for everything else: 3-D tensors with a different/static leading index, broadcast variants @@ -370,6 +376,18 @@ def aux_tensor_load_kind( if len(carrier_tile_shape) != 2: return None + if explicit_rank1_data_axis is not None: + return _matches_rank1_broadcast( + aux_shape=aux_shape, + aux_tensor_shape=aux_tensor_shape, + aux_tensor_val=aux_tensor_val, + index_list=index_list, + carrier_tile_shape=carrier_tile_shape, + carrier_tile_index_nodes=carrier_tile_index_nodes, + carrier_global_shape=carrier_global_shape, + data_axis=explicit_rank1_data_axis, + ) + colvec = _matches_colvec_broadcast( aux_shape=aux_shape, aux_tensor_shape=aux_tensor_shape, @@ -434,10 +452,9 @@ def _matches_broadcast( aligns to the last dimension under PyTorch broadcasting rules, so ``acc + bias[tile_m]`` is not a column-vector broadcast — it is a shape error (BM ≠ BN) or trailing-axis broadcast (BM == BN). - Users wanting an explicit M-axis broadcast must write it - explicitly (``bias[tile_m][:, None]`` / ``.unsqueeze(-1)`` / - ``.expand(...)``); that is a separate pattern handler not yet - wired up. + Users wanting an M-axis broadcast must write it explicitly as + ``bias[tile_m].unsqueeze(-1)``; the operand normalizer supplies axis 0 + to :func:`_matches_rank1_broadcast` for that form. The splice site emits ``cute.make_layout((m, n), stride=(0, 1))`` with stride 1 hard-coded on the data axis, @@ -455,42 +472,54 @@ def _matches_broadcast( to be divisible by the tile extent passes the tile check yet reads OOB at runtime. """ - if len(aux_tensor_shape) != 1 or len(aux_shape) != 1 or len(index_list) != 1: - return None - if carrier_tile_index_nodes is None: - # Broadcast classification is gated on knowing the carrier's - # tile-id symbols so we can pin which axis the load broadcasts - # along. Without that, the unpredicated splice could broadcast - # the wrong axis silently — bail to the loud-failure backstop. - return None - if len(carrier_tile_index_nodes) != 2: + return _matches_rank1_broadcast( + aux_shape=aux_shape, + aux_tensor_shape=aux_tensor_shape, + aux_tensor_val=aux_tensor_val, + index_list=index_list, + carrier_tile_shape=carrier_tile_shape, + carrier_tile_index_nodes=carrier_tile_index_nodes, + carrier_global_shape=carrier_global_shape, + data_axis=1, + ) + + +def _matches_rank1_broadcast( + *, + aux_shape: tuple[object, ...], + aux_tensor_shape: tuple[object, ...], + aux_tensor_val: torch.Tensor, + index_list: Sequence[object], + carrier_tile_shape: tuple[object, ...], + carrier_tile_index_nodes: tuple[torch.fx.Node, ...] | None, + carrier_global_shape: tuple[object, ...] | None, + data_axis: int, +) -> tuple[str, int] | None: + """Match a contiguous rank-1 load against one carrier data axis.""" + if ( + data_axis not in (0, 1) + or len(aux_tensor_shape) != 1 + or len(aux_shape) != 1 + or len(index_list) != 1 + or carrier_tile_index_nodes is None + or len(carrier_tile_index_nodes) != 2 + or aux_tensor_val.stride() != (1,) + ): return None load_idx = index_list[0] - if not isinstance(load_idx, torch.fx.Node): - return None - if aux_tensor_val.stride() != (1,): - return None - aux_extent = aux_shape[0] - # Trailing-axis (rowvec) broadcast only — see docstring. - axis = 1 - if load_idx is not carrier_tile_index_nodes[axis]: + if ( + load_idx is not carrier_tile_index_nodes[data_axis] + or aux_shape[0] != carrier_tile_shape[data_axis] + ): return None - if aux_extent != carrier_tile_shape[axis]: + if carrier_global_shape is not None and ( + len(carrier_global_shape) != 2 + or aux_tensor_shape[0] != carrier_global_shape[data_axis] + ): return None - if carrier_global_shape is not None: - if len(carrier_global_shape) != 2: - return None - # The underlying rank-1 aux tensor's *total* length must - # match the output's global axis size on the broadcast - # axis. The earlier ``aux_extent == carrier_tile_shape`` - # check matches the per-load *tile* extent against the - # carrier's tile shape; without the global-size check, an - # aux smaller than the output axis whose length happens to - # be divisible by the tile extent could pass classification - # yet read OOB at runtime. - if aux_tensor_shape[0] != carrier_global_shape[axis]: - return None - return ("broadcast", axis) + # The existing column-vector path is identified by 2; row vectors retain + # their historical axis-1 representation. + return ("broadcast", 2 if data_axis == 0 else 1) def _matches_leading_broadcast( diff --git a/helion/language/memory_ops.py b/helion/language/memory_ops.py index 01e37cc18..d62daefed 100644 --- a/helion/language/memory_ops.py +++ b/helion/language/memory_ops.py @@ -1744,12 +1744,17 @@ def _codegen_cute_store_tcgen05_tile( # treats unreferenced tensors as captures, which doesn't work # for tensors only read inside a per-subtile loop body). df.placeholder_args.add(aux_tensor_name) - # Broadcast axes 0/1 build a stride-0 2-D view of a rank-1 tensor. + # Compact broadcast inputs build a logical 2-D view over their rank-1 + # storage. Full-shape colvec views already carry stride (1, 0). # An exact rank-3 input uses a view of its selected leading coordinate. - # The colvec form (2) already carries an (M, N) stride-(1, 0) view. + compact_broadcast = ( + aux_step.broadcast_axis is not None and aux_torch_tensor.ndim == 1 + ) aux_view2d = ( df.new_var(f"tcgen05_aux_view2d_{aux_idx}") - if aux_step.broadcast_axis in (0, 1) or aux_torch_tensor.ndim == 3 + if compact_broadcast + or aux_step.broadcast_axis == 0 + or aux_torch_tensor.ndim == 3 else None ) aux_step_records.append( @@ -2268,7 +2273,21 @@ def _aux_tile_setup_lines( ) continue - if rec.broadcast_axis is None or rec.broadcast_axis == 2: + compact_colvec = rec.broadcast_axis == 2 and rec.aux_view2d is not None + if compact_colvec: + assert rec.aux_view2d is not None + # Explicit ``rank1.unsqueeze(-1)`` column vector. Reconstruct + # its logical (M, N) stride-(1, 0) view without requiring the + # caller to materialize an expanded host tensor. + lines.append( + f"{rec.aux_view2d} = cute.make_tensor(" + f"{rec.aux_tensor_name}.iterator, " + f"cute.make_layout(({m_size}, {n_size}), " + f"stride=(1, 0)))" + ) + source_for_local_tile = rec.aux_view2d + aux_tile_is_local = False + elif rec.broadcast_axis is None or rec.broadcast_axis == 2: # Exact-shape aux (or the colvec form) uses the trailing matrix # directly. For a rank-3 residual, first select the current # leading-passthrough slice without changing its M/N strides. diff --git a/test/test_cute_backend.py b/test/test_cute_backend.py index 248be68c3..99074c6e7 100644 --- a/test/test_cute_backend.py +++ b/test/test_cute_backend.py @@ -1056,7 +1056,9 @@ def cute_rhs_batched_dot_bias_tcgen05( acc = hl.zeros([tile_b, tile_m, tile_n], dtype=torch.float32) for tile_k in hl.tile(k): acc = hl.dot(w[tile_m, tile_k], y[tile_b, tile_k, tile_n], acc=acc) - out[tile_b, tile_m, tile_n] = (acc + bias[tile_n]).to(torch.bfloat16) + out[tile_b, tile_m, tile_n] = (acc + bias[tile_m].unsqueeze(-1)).to( + torch.bfloat16 + ) return out @@ -5399,7 +5401,7 @@ def test_one_cta_k_tail_coverage_is_generic(self) -> None: [1, 128, 128, 64], "persistent_interleaved", torch.matmul(rhs_leading_args[0].float(), rhs_leading_args[1].float()) - + rhs_leading_args[2].float(), + + rhs_leading_args[2].float()[None, :, None], "'rhs_leading_passthrough': True", ), ) diff --git a/test/test_cute_lowerings.py b/test/test_cute_lowerings.py index 74e44fefd..80a366b1d 100644 --- a/test/test_cute_lowerings.py +++ b/test/test_cute_lowerings.py @@ -6836,9 +6836,8 @@ def test_tcgen05_fused_colvec_broadcast_rejected_at_classify_time( ``carrier_tile_index_nodes[1]`` (the trailing axis). Users wanting an explicit colvec broadcast must spell it - out (``bias[tile_m][:, None]`` / ``.unsqueeze(-1)`` / - ``.expand(...)``); that is a separate, deferred pattern - handler. + with ``.unsqueeze(-1)``; the adjacent positive test covers + that distinct, unambiguous expression. """ from helion._compiler.cute.mma_support import get_cute_mma_support @@ -6883,6 +6882,42 @@ def cute_matmul_colvec( self.assertIn("tcgen05 MMA path", message) self.assertIn("rowvec", message) + def test_tcgen05_fused_rank1_colvec_unsqueeze_runtime_correctness(self) -> None: + """An explicit rank-1 ``unsqueeze(-1)`` uses the colvec epilogue.""" + + from helion._compiler.cute.mma_support import get_cute_mma_support + + if not get_cute_mma_support().tcgen05_f16bf16: + self.skipTest("tcgen05 F16/BF16 MMA is not supported on this machine") + + @helion.kernel(backend="cute") + def cute_matmul_colvec( + x: torch.Tensor, y: torch.Tensor, colvec: torch.Tensor + ) -> torch.Tensor: + m, k = x.size() + _, n = y.size() + out = torch.empty([m, n], dtype=x.dtype, device=x.device) + for tile_m, tile_n in hl.tile([m, n]): + acc = hl.zeros([tile_m, tile_n], dtype=torch.float32) + for tile_k in hl.tile(k): + acc = torch.addmm(acc, x[tile_m, tile_k], y[tile_k, tile_n]) + out[tile_m, tile_n] = (acc + colvec[tile_m].unsqueeze(-1)).to(x.dtype) + return out + + x = torch.randn(128, 64, dtype=torch.bfloat16, device=DEVICE) + y = torch.randn(64, 64, dtype=torch.bfloat16, device=DEVICE) + colvec = torch.randn(128, dtype=torch.bfloat16, device=DEVICE) + bound = cute_matmul_colvec.bind((x, y, colvec)) + config = _make_tcgen05_persistent_config( + block_sizes=[128, 64, 32], + pid_type="persistent_interleaved", + ) + code = bound.to_triton_code(config) + out = bound.compile_config(config)(x, y, colvec) + expected = (x.float() @ y.float() + colvec.float()[:, None]).to(x.dtype) + torch.testing.assert_close(out, expected, atol=2e-1, rtol=1e-2) + self.assertIn("stride=(1, 0)", code) + def test_tcgen05_fused_bias_broadcast_codegen_marker(self) -> None: """The spliced rowvec broadcast emits an ``aux_view2d`` local bound to ``cute.make_tensor(bias.iterator, @@ -14127,6 +14162,29 @@ def test_aux_load_kind_colvec_stride_1_0_is_broadcast_2(self) -> None: ("broadcast", 2), ) + def test_aux_load_kind_explicit_rank1_broadcast_axes(self) -> None: + from helion._compiler.cute.cute_fx_walk import aux_tensor_load_kind + + idx = self._carrier_index_nodes() + vector = torch.empty(4, dtype=torch.float32) + for data_axis, expected_kind in ((0, ("broadcast", 2)), (1, ("broadcast", 1))): + with self.subTest(data_axis=data_axis): + load_node, _ = self._build_aux_load_node( + aux_tensor=vector, + load_shape=(4,), + index_nodes=(idx[data_axis],), + ) + self.assertEqual( + aux_tensor_load_kind( + load_node, + carrier_tile_shape=(4, 4), + carrier_tile_index_nodes=idx, + carrier_global_shape=(4, 4), + explicit_rank1_data_axis=data_axis, + ), + expected_kind, + ) + def test_aux_load_kind_exact_tensor_is_not_colvec(self) -> None: """A genuine dense ``(M, N)`` aux (trailing stride 1) must fall through the colvec matcher to ``("exact", None)`` — the colvec