| |
| """C28: the gather-GEMM sparse-conv encoder on the device (PLAN.md 1.2 row C28; probe P4; host rulebooks: C13 |
| ``ttaw.sparse``). Owned by the BEVFusion port; PTv3's sub-manifold stages are the second consumer. |
| |
| Every spconv layer of the Autoware graphs (sub-manifold and strided alike) is a gather (C13): ``y[o] = sum_k |
| W[:, k, :] @ x[nmap[o, k]] + b``. On the device one layer is: |
| |
| 1. ONE ``ttnn.embedding`` (TILE + PADDED, P4) of the layer's tile-ordered im2col index (C13 ``im2col_index``, |
| uint32 ``[rows / 32, K * 32]``) into the bf16 ROW_MAJOR table ``[1, 1, V + 1, Cin_pad]`` of the layer's input rows |
| (the last row = the explicit zero sentinel ``V``, ``Cin_pad % 32 == 0``) -> ``[rows / 32, K * 32, Cin_pad]`` |
| whose tiles, in memory order, are the tiles of the im2col matrix ``[rows, K * Cin_pad]``; |
| 2. ``ttnn.experimental.view`` to ``[1, 1, rows, K * Cin_pad]`` (metadata only); |
| 3. a bias-free ``ttnn.matmul`` with the BN-folded im2col weight ``[K * Cin_pad, Cout_pad]`` into fp32 (C24 |
| ``RowMatmul``: no fused-bias epilogue rounds the fp32 accumulator), optionally a second one with the weight's |
| bf16 low term (``weight_terms=2``); |
| 4. one exact fp32 ``ttnn.add`` of the bias (with the ReLU fused when there is no residual), or bias + the fp32 |
| identity of a basic block followed by the ReLU fused into the second add; |
| 5. the next table: the fp32 output cast to bf16, untilized and given its zero row (C19 ``with_zero_rows``). |
| |
| **Activation terms.** ``ttnn.embedding`` is bf16-only, so a table rounds every layer's input to bf16 (8 significant |
| bits). A layer whose table is *two-term* (``in_terms=2``) carries ``hi = bf16(x)`` in channels ``[0, C)`` and ``lo = |
| bf16(x - hi)`` in ``[C, 2C)`` of the same table (~16 significant bits): the im2col weight repeats each tap's rows for |
| the ``lo`` block, so the same single gather and single matmul compute ``hi @ W + lo @ W``. Where ``2C <= 32`` (Cin 4 |
| of a VFE input, 16 at BEVFusion's level 1) the ``lo`` block rides in channels that the tile padding carries anyway: |
| free. A producer with ``out_terms=2`` writes its output columns twice (``[W | W]``) and splits them on the device into |
| ``[hi | lo]`` (``y - bf16(y) * mask``, then one bf16 cast: exact). Consecutive layers must agree (the producer's |
| ``out_terms`` is the consumer's ``in_terms``; :class:`SparseEncoder` checks it). |
| |
| Rows past a level's active count are computed (capacity padding: nothing data-dependent reaches the device) but |
| never gathered: every index points at an active row or at the sentinel, so no count is needed on the device. |
| |
| numpy only at import (ttnn / torch inside the device calls). Host oracle: :func:`sparse_encoder_numpy`, with the |
| device's bf16 table rounding modelled (``round_tables=True``). |
| """ |
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple |
|
|
| import numpy as np |
|
|
| from ..precision import Precision, PrecisionPolicy |
| from .second import free_maps |
|
|
| __all__ = ["SparseConvSpec", "TtSparseConv", "SparseEncoder", "DEFAULT_PRECISION", "TILE_ROWS", "pad_channels", |
| "pack_table", "table_rows_numpy", "sparse_layer_numpy", "sparse_encoder_numpy", "dense_gather", |
| "im2col_weight"] |
|
|
| TILE_ROWS = 32 |
| |
| DEFAULT_PRECISION = "HiFi4+fp32:w=fp32" |
|
|
|
|
| def pad_channels(c: int, multiple: int = TILE_ROWS) -> int: |
| """``c`` rounded up to a multiple of 32 (the table width a TILE embedding needs, P4).""" |
| return -(-int(c) // multiple) * multiple |
|
|
|
|
| def _bf16(a: np.ndarray) -> np.ndarray: |
| from ..tensors import round_to_bf16 |
|
|
| return round_to_bf16(np.asarray(a, np.float32)) |
|
|
|
|
| @dataclass(frozen=True) |
| class SparseConvSpec: |
| """One BN-folded sparse conv: ``y = relu?(gather_conv(x, weight) + bias (+ identity))``. |
| |
| ``weight``: KRSC ``[Cout, k0, k1, k2, Cin]`` float32 (tap ``k = (a * k1 + b) * k2 + c``, C13 order) or already |
| ``[Cout, K, Cin]``; ``index``: the name of the gather index it reads (several layers may share one: the |
| sub-manifold convs of a level); ``residual``: the name of an earlier layer whose (fp32, pre-split) output is |
| added before the ReLU; ``in_terms`` / ``out_terms``: 1 (bf16 table) or 2 (``[hi | lo]`` table, module |
| docstring).""" |
|
|
| name: str |
| weight: np.ndarray |
| bias: np.ndarray |
| index: str |
| residual: Optional[str] = None |
| relu: bool = True |
| in_terms: int = 1 |
| out_terms: int = 1 |
|
|
| def __post_init__(self) -> None: |
| w = np.asarray(self.weight, np.float32) |
| if w.ndim not in (3, 5): |
| raise ValueError(f"{self.name}: weight must be KRSC [Cout, k0, k1, k2, Cin] or [Cout, K, Cin], got {w.shape}") |
| if np.shape(self.bias) != (w.shape[0],): |
| raise ValueError(f"{self.name}: bias must be ({w.shape[0]},), got {np.shape(self.bias)}") |
| if self.in_terms not in (1, 2) or self.out_terms not in (1, 2): |
| raise ValueError(f"{self.name}: terms must be 1 or 2") |
| if not (np.all(np.isfinite(w)) and np.all(np.isfinite(self.bias))): |
| raise ValueError(f"{self.name}: non-finite weights") |
|
|
| @property |
| def weight3(self) -> np.ndarray: |
| w = np.asarray(self.weight, np.float32) |
| return w.reshape(w.shape[0], -1, w.shape[-1]) |
|
|
| @property |
| def kernel_volume(self) -> int: |
| return int(self.weight3.shape[1]) |
|
|
| @property |
| def in_channels(self) -> int: |
| return int(self.weight3.shape[2]) |
|
|
| @property |
| def out_channels(self) -> int: |
| return int(self.weight3.shape[0]) |
|
|
| @property |
| def cin_pad(self) -> int: |
| """Table width of the input: ``in_terms * Cin`` rounded up to 32.""" |
| return pad_channels(self.in_terms * self.in_channels) |
|
|
| @property |
| def cout_pad(self) -> int: |
| """Output / next-table width: ``out_terms * Cout`` rounded up to 32.""" |
| return pad_channels(self.out_terms * self.out_channels) |
|
|
| def matrix(self) -> np.ndarray: |
| """The im2col weight ``[K * cin_pad, cout_pad]`` float32 (:func:`im2col_weight`).""" |
| return im2col_weight(self.weight3, cin_pad=self.cin_pad, cout_pad=self.cout_pad, in_terms=self.in_terms, |
| out_terms=self.out_terms) |
|
|
| def bias_row(self) -> np.ndarray: |
| """``(1, cout_pad)`` float32: the bias, repeated for the ``lo`` block when ``out_terms == 2``.""" |
| out = np.zeros((1, self.cout_pad), np.float32) |
| c = self.out_channels |
| for t in range(self.out_terms): |
| out[0, t * c:(t + 1) * c] = np.asarray(self.bias, np.float32) |
| return out |
|
|
| def lo_mask(self) -> np.ndarray: |
| """``(1, cout_pad)`` float32: 1 on the ``lo`` output columns (``out_terms == 2``), else 0.""" |
| out = np.zeros((1, self.cout_pad), np.float32) |
| if self.out_terms == 2: |
| out[0, self.out_channels:2 * self.out_channels] = 1.0 |
| return out |
|
|
|
|
| def im2col_weight(weight3: np.ndarray, *, cin_pad: int, cout_pad: int, in_terms: int = 1, |
| out_terms: int = 1) -> np.ndarray: |
| """``[Cout, K, Cin]`` -> ``[K * cin_pad, cout_pad]``: row ``k * cin_pad + j`` = ``W[:, k, j mod Cin]`` for the |
| ``in_terms`` channel blocks (the ``lo`` block repeats the ``hi`` weights), zero for padding; columns repeated for |
| the ``out_terms`` blocks, zero padding beyond.""" |
| w = np.asarray(weight3, np.float32) |
| cout, K, cin = w.shape |
| if in_terms * cin > cin_pad or out_terms * cout > cout_pad: |
| raise ValueError(f"cin_pad {cin_pad} / cout_pad {cout_pad} too small for {in_terms} x {cin} / " |
| f"{out_terms} x {cout}") |
| m = np.zeros((K, cin_pad, cout_pad), np.float32) |
| wt = np.transpose(w, (1, 2, 0)) |
| for ti in range(in_terms): |
| for to in range(out_terms): |
| m[:, ti * cin:(ti + 1) * cin, to * cout:(to + 1) * cout] = wt |
| return m.reshape(K * cin_pad, cout_pad) |
|
|
|
|
| |
|
|
| def table_rows_numpy(x: np.ndarray, *, terms: int, width: int) -> np.ndarray: |
| """``(N, C)`` float rows -> ``(N, width)`` float32 table rows with the device's bf16 values: ``terms=1``: |
| ``bf16(x)`` in ``[0, C)``; ``terms=2``: ``hi = bf16(x)`` in ``[0, C)``, ``lo = bf16(x - hi)`` in ``[C, 2C)``; |
| zero elsewhere.""" |
| x = np.asarray(x, np.float32) |
| n, c = x.shape |
| if terms * c > width: |
| raise ValueError(f"{terms} x {c} channels do not fit a width of {width}") |
| out = np.zeros((n, width), np.float32) |
| hi = _bf16(x) |
| out[:, :c] = hi |
| if terms == 2: |
| out[:, c:2 * c] = _bf16(x - hi) |
| return out |
|
|
|
|
| def pack_table(x: np.ndarray, capacity: int, *, terms: int = 1, width: Optional[int] = None, |
| out: Optional[np.ndarray] = None) -> np.ndarray: |
| """The first layer's input table as a host array ``(1, 1, capacity + 1, width)`` float32 (bf16-exact values, |
| upload as bf16 ROW_MAJOR): the ``N <= capacity`` rows of ``x`` (:func:`table_rows_numpy`), zero rows after them, |
| and the zero sentinel row at ``capacity``. ``width`` defaults to ``terms * C`` rounded up to 32.""" |
| x = np.asarray(x, np.float32) |
| n, c = x.shape |
| if n > capacity: |
| raise ValueError(f"{n} rows exceed the capacity {capacity}") |
| width = pad_channels(terms * c) if width is None else int(width) |
| shape = (1, 1, int(capacity) + 1, width) |
| if out is None: |
| out = np.zeros(shape, np.float32) |
| elif out.shape != shape or out.dtype != np.float32: |
| raise ValueError(f"out must be float32 {shape}") |
| else: |
| out[...] = 0.0 |
| out[0, 0, :n] = table_rows_numpy(x, terms=terms, width=width) |
| return out |
|
|
|
|
| def sparse_layer_numpy(table: np.ndarray, nmap: np.ndarray, spec: SparseConvSpec, *, |
| identity: Optional[np.ndarray] = None) -> np.ndarray: |
| """Host oracle of one device layer on a host table ``(V + 1, cin_pad)`` (the last row the zero sentinel): |
| gathers by ``nmap`` (``-1`` -> sentinel), the im2col matmul in float64, bias, identity, ReLU -> ``(N, Cout)`` |
| float32 (the logical columns).""" |
| t = np.asarray(table, np.float64).reshape(-1, spec.cin_pad) |
| sentinel = t.shape[0] - 1 |
| idx = np.where(np.asarray(nmap) >= 0, nmap, sentinel).astype(np.int64) |
| cols = t[idx].reshape(idx.shape[0], -1) |
| m = spec.matrix().astype(np.float64)[:, :spec.out_channels] |
| y = cols @ m + np.asarray(spec.bias, np.float64) |
| if identity is not None: |
| y = y + np.asarray(identity, np.float64) |
| if spec.relu: |
| y = np.maximum(y, 0.0) |
| return y.astype(np.float32) |
|
|
|
|
| def sparse_encoder_numpy(specs: Sequence[SparseConvSpec], features: np.ndarray, maps: Mapping[str, np.ndarray], *, |
| round_tables: bool = True) -> Dict[str, np.ndarray]: |
| """Host oracle of :class:`SparseEncoder`: ``features`` ``(N0, C0)`` (the first layer's input rows) and the |
| neighbour maps by index name -> ``{layer name: (N, Cout) float32}``. ``round_tables=True`` models the device's |
| tables (bf16 or two-term ``hi + lo``, per the specs' terms); False is the fp32 network (the reference).""" |
| outs: Dict[str, np.ndarray] = {} |
| x = np.asarray(features, np.float32) |
| for spec in specs: |
| if round_tables: |
| table = table_rows_numpy(x, terms=spec.in_terms, width=spec.cin_pad) |
| else: |
| table = np.zeros((x.shape[0], spec.cin_pad), np.float32) |
| table[:, :x.shape[1]] = x |
| table = np.concatenate([table, np.zeros((1, spec.cin_pad), np.float32)], 0) |
| ident = outs[spec.residual] if spec.residual else None |
| y = sparse_layer_numpy(table, maps[spec.index], spec, identity=ident) |
| outs[spec.name] = y |
| x = y |
| return outs |
|
|
|
|
| |
|
|
| class TtSparseConv: |
| """One device layer of :class:`SparseConvSpec` (module docstring, steps 1-4). ``weight_terms=2`` runs the |
| im2col weight as two bf16 terms (two matmuls, ~16 significant bits) instead of fp32 weights (TF32-like |
| products, P12). Device weights are uploaded by the first call (outside any capture).""" |
|
|
| def __init__(self, spec: SparseConvSpec, *, precision: Any = DEFAULT_PRECISION, weight_terms: int = 1): |
| from .pillars import RowMatmul |
|
|
| self.spec = spec |
| self.name = spec.name |
| self.precision = Precision.parse(precision) if not isinstance(precision, Precision) else precision |
| self.weight_terms = int(weight_terms) |
| if self.weight_terms not in (1, 2): |
| raise ValueError(f"{self.name}: weight_terms must be 1 or 2") |
| m = spec.matrix() |
| if self.weight_terms == 2: |
| if self.precision.weights != "bfloat16": |
| self.precision = self.precision.with_(weights="bfloat16") |
| hi = _bf16(m) |
| self.mm = [RowMatmul(hi, precision=self.precision, name=f"{self.name}.whi"), |
| RowMatmul(_bf16(m - hi), precision=self.precision, name=f"{self.name}.wlo")] |
| else: |
| self.mm = [RowMatmul(m, precision=self.precision, name=self.name)] |
| self.k_in = int(m.shape[0]) |
| self._consts: Optional[Tuple[Any, Any]] = None |
|
|
| @property |
| def rows_per_block(self) -> int: |
| return TILE_ROWS |
|
|
| def _const(self, device: Any) -> Tuple[Any, Any]: |
| import torch |
| import ttnn |
|
|
| if self._consts is None: |
| def up(a): |
| return ttnn.from_torch(torch.from_numpy(np.ascontiguousarray(a, np.float32)), dtype=ttnn.float32, |
| layout=ttnn.TILE_LAYOUT, device=device, memory_config=ttnn.DRAM_MEMORY_CONFIG) |
|
|
| mask = up(self.spec.lo_mask()) if self.spec.out_terms == 2 else None |
| self._consts = (up(self.spec.bias_row()), mask) |
| return self._consts |
|
|
| def gather(self, table: Any, index: Any): |
| """The im2col operand ``[1, 1, rows, K * cin_pad]`` bf16 TILE (one PADDED embedding + a view).""" |
| import ttnn |
|
|
| width = int(table.shape[-1]) |
| if width != self.spec.cin_pad: |
| raise ValueError(f"{self.name}: table width {width} != cin_pad {self.spec.cin_pad}") |
| k = self.spec.kernel_volume |
| if int(index.shape[-1]) != k * TILE_ROWS: |
| raise ValueError(f"{self.name}: index has {int(index.shape[-1])} columns, expected {k * TILE_ROWS}") |
| rows = int(index.shape[0]) * TILE_ROWS |
| sentinel = int(table.shape[-2]) - 1 |
| g = ttnn.embedding(index, table, layout=ttnn.TILE_LAYOUT, padding_idx=sentinel, |
| embeddings_type=ttnn.EmbeddingsType.PADDED) |
| return ttnn.experimental.view(g, [1, 1, rows, k * width]) |
|
|
| def __call__(self, table: Any, index: Any, identity: Any = None): |
| """``table`` bf16 ROW_MAJOR ``[1, 1, V + 1, cin_pad]`` (row V zero), ``index`` uint32 ``[rows / 32, K * 32]``, |
| ``identity`` an fp32 TILE ``[1, 1, rows, cout_pad]`` (the residual source's output) or None -> the fp32 TILE |
| output ``[1, 1, rows, cout_pad]`` (DRAM). The inputs are kept.""" |
| import ttnn |
|
|
| from ..ops.conv import unary_with_param |
|
|
| x = self.gather(table, index) |
| ys = [mm(x) for mm in self.mm] |
| free_maps(x) |
| y = ys[-1] |
| for extra in reversed(ys[:-1]): |
| s = ttnn.add(extra, y, memory_config=ttnn.DRAM_MEMORY_CONFIG) |
| free_maps(extra, y) |
| y = s |
| bias, _ = self._const(table.device()) |
| relu = [unary_with_param("relu")] if self.spec.relu else [] |
| if identity is None: |
| out = ttnn.add(y, bias, activations=relu, memory_config=ttnn.DRAM_MEMORY_CONFIG) |
| free_maps(y) |
| return out |
| s = ttnn.add(y, bias, memory_config=ttnn.DRAM_MEMORY_CONFIG) |
| free_maps(y) |
| out = ttnn.add(s, identity, activations=relu, memory_config=ttnn.DRAM_MEMORY_CONFIG) |
| free_maps(s) |
| return out |
|
|
| def to_table(self, y: Any): |
| """The fp32 output -> the next layer's bf16 ROW_MAJOR table ``[1, 1, rows + 1, cout_pad]`` (zero last row; |
| ``[hi | lo]`` split for ``out_terms == 2``). ``y`` is kept.""" |
| import ttnn |
|
|
| from ..ops.gather import with_zero_rows |
|
|
| if self.spec.out_terms == 2: |
| _, mask = self._const(y.device()) |
| hi = ttnn.typecast(y, ttnn.bfloat16) |
| h32 = ttnn.typecast(hi, ttnn.float32) |
| free_maps(hi) |
| hm = ttnn.multiply(h32, mask) |
| free_maps(h32) |
| r = ttnn.subtract(y, hm) |
| free_maps(hm) |
| t = ttnn.typecast(r, ttnn.bfloat16) |
| free_maps(r) |
| else: |
| t = ttnn.typecast(y, ttnn.bfloat16) |
| table = with_zero_rows(t, 1) |
| free_maps(t) |
| return table |
|
|
| def release(self) -> None: |
| for mm in self.mm: |
| mm.release() |
| if self._consts is not None: |
| free_maps(*[c for c in self._consts if c is not None]) |
| self._consts = None |
|
|
| def describe(self) -> Dict[str, Any]: |
| s = self.spec |
| return {"name": s.name, "index": s.index, "cin": s.in_channels, "cout": s.out_channels, "K": s.kernel_volume, |
| "cin_pad": s.cin_pad, "cout_pad": s.cout_pad, "in_terms": s.in_terms, "out_terms": s.out_terms, |
| "residual": s.residual, "relu": s.relu, "weight_terms": self.weight_terms, |
| "precision": self.precision.label} |
|
|
|
|
| class SparseEncoder: |
| """A chain of :class:`TtSparseConv` layers (in order; each reads the previous layer's table, the first one the |
| given input table). ``policy`` resolves ``<prefix-free layer name>`` to a precision (default |
| :data:`DEFAULT_PRECISION`); ``weight_terms`` 1 or 2 for every layer (or a ``{name: terms}`` mapping).""" |
|
|
| def __init__(self, specs: Sequence[SparseConvSpec], *, policy: Optional[PrecisionPolicy] = None, |
| weight_terms: Any = 1, name: str = "sparse"): |
| if not specs: |
| raise ValueError("no layers") |
| self.name = name |
| names = [s.name for s in specs] |
| if len(set(names)) != len(names): |
| raise ValueError("layer names must be unique") |
| for prev, cur in zip(specs[:-1], specs[1:]): |
| if prev.out_terms != cur.in_terms or prev.cout_pad != cur.cin_pad: |
| raise ValueError(f"{cur.name}: input terms / width {cur.in_terms} / {cur.cin_pad} != the output of " |
| f"{prev.name} ({prev.out_terms} / {prev.cout_pad})") |
| if prev.out_channels != cur.in_channels: |
| raise ValueError(f"{cur.name}: {cur.in_channels} input channels after {prev.out_channels}") |
| seen = set() |
| for s in specs: |
| if s.residual is not None: |
| if s.residual not in seen: |
| raise ValueError(f"{s.name}: residual {s.residual!r} is not an earlier layer") |
| src = next(p for p in specs if p.name == s.residual) |
| if src.cout_pad != s.cout_pad or src.out_channels != s.out_channels or src.out_terms != s.out_terms: |
| raise ValueError(f"{s.name}: residual {s.residual!r} has another output layout") |
| seen.add(s.name) |
| self.specs = list(specs) |
| self.layers: List[TtSparseConv] = [] |
| for s in specs: |
| prec = policy.resolve(s.name) if policy is not None else Precision.parse(DEFAULT_PRECISION) |
| wt = int(weight_terms.get(s.name, 1)) if isinstance(weight_terms, Mapping) else int(weight_terms) |
| self.layers.append(TtSparseConv(s, precision=prec, weight_terms=wt)) |
| |
| self._last_use: Dict[str, int] = {} |
| for i, s in enumerate(specs): |
| if s.residual is not None: |
| self._last_use[s.residual] = i |
|
|
| @property |
| def index_names(self) -> List[str]: |
| """The gather indices the chain reads, in first-use order.""" |
| out: List[str] = [] |
| for s in self.specs: |
| if s.index not in out: |
| out.append(s.index) |
| return out |
|
|
| @property |
| def input_width(self) -> int: |
| return self.specs[0].cin_pad |
|
|
| @property |
| def input_terms(self) -> int: |
| return self.specs[0].in_terms |
|
|
| def __call__(self, table: Any, indices: Mapping[str, Any], taps: Optional[Dict[str, Any]] = None, |
| keep_table: bool = True): |
| """``table``: the first layer's bf16 ROW_MAJOR input table (kept); ``indices``: ``{index name: uint32 device |
| tensor}``; ``taps`` (a dict) collects every layer's fp32 output ``[1, 1, rows, cout_pad]`` by layer name. |
| Returns ``(y, out_table)``: the last layer's fp32 output and (``keep_table``) its bf16 ROW_MAJOR table |
| (zero sentinel row appended), else ``(y, None)``.""" |
| missing = [n for n in self.index_names if n not in indices] |
| if missing: |
| raise KeyError(f"{self.name}: missing gather indices {missing}") |
| cur_table = table |
| outs: Dict[str, Any] = {} |
| y = None |
| n = len(self.layers) |
| for i, (layer, spec) in enumerate(zip(self.layers, self.specs)): |
| ident = outs[spec.residual] if spec.residual else None |
| y = layer(cur_table, indices[spec.index], ident) |
| if cur_table is not table: |
| free_maps(cur_table) |
| |
| for name, last in list(self._last_use.items()): |
| if last == i and name in outs and (taps is None): |
| free_maps(outs.pop(name)) |
| outs[spec.name] = y |
| if taps is not None: |
| taps[spec.name] = y |
| last_layer = i == n - 1 |
| if not last_layer or keep_table: |
| cur_table = layer.to_table(y) |
| if not last_layer and taps is None and spec.name not in self._last_use: |
| free_maps(outs.pop(spec.name)) |
| return y, (cur_table if keep_table else None) |
|
|
| def release(self) -> None: |
| for layer in self.layers: |
| layer.release() |
|
|
| def describe(self) -> Dict[str, Any]: |
| return {"name": self.name, "layers": [layer.describe() for layer in self.layers], |
| "indices": self.index_names} |
|
|
|
|
| def dense_gather(table: Any, index: Any): |
| """To-dense as a gather (C13 ``dense_gather_index``, one slice): ``table`` bf16 ROW_MAJOR ``[1, 1, V + 1, C]`` |
| (row V zero), ``index`` uint32 ``[1, cells]`` -> ``[1, 1, cells, C]`` bf16 TILE (empty cells exactly 0).""" |
| from ..ops.gather import gather_rows |
|
|
| return gather_rows(index, table, sentinel=int(table.shape[-2]) - 1) |
|
|
