| |
| """Host packing of one plan's :class:`host.pipeline.Prepared` into the persistent trace inputs (numpy, no ttnn). |
| |
| Every array is float32 in the 4-D shape of its device input (TILE, uploaded each plan): the mixer inputs keep only |
| the time rows the export reads (ego 0..5, neighbours 25..30: the pad-relative island needs nothing else), the aux |
| columns of the exact embedding rewrites (``tt/params.py``), the 576-row token arrays (564 real + 12 zero pad |
| tokens), the key-bias rows (0 / -inf; extra keys masked) and the solver's ``cs`` / ``y0`` on 352 rows. |
| """ |
| from __future__ import annotations |
|
|
| from typing import Dict, Optional |
|
|
| import numpy as np |
|
|
| from ..host.pipeline import Prepared |
| from ..reference import config as C |
| from ..ttaw.ops.attention import key_bias_row |
| from . import config as T |
| from .params import lane_aux_features, pad_rows |
|
|
| __all__ = ["plan_inputs", "INPUT_SPECS", "warmup_inputs", "decoder_state", "BF16_INPUTS", "CS_COLS"] |
|
|
| f32 = np.float32 |
|
|
| |
| CS_COLS = T.TILE if T.KNOBS.read().INPUT_TRIM else T.STATE_COLS |
|
|
| |
| INPUT_SPECS: Dict[str, tuple] = { |
| "ego_x": (1, 1, T.MIXER_T["ego"], T.MIXER_CIN["ego"]), |
| "neighbor_x": (1, C.MAX_NUM_NEIGHBORS, T.MIXER_T["neighbor"], T.MIXER_CIN["neighbor"]), |
| "neighbor_aux": (1, 1, C.MAX_NUM_NEIGHBORS, T.NEIGHBOR_AUX_DIM), |
| "static_x": (1, 1, C.NUM_STATIC_OBJECTS, C.STATIC_OBJECT_DIM), |
| "lane_x": (1, C.NUM_SEGMENTS_IN_LANE, T.MIXER_T["lane"], T.MIXER_CIN["lane"]), |
| "lane_aux": (1, 1, C.NUM_SEGMENTS_IN_LANE, T.LANE_AUX_DIM), |
| "route_x": (1, C.NUM_SEGMENTS_IN_ROUTE, T.MIXER_T["route"], T.MIXER_CIN["route"]), |
| "route_aux": (1, 1, C.NUM_SEGMENTS_IN_ROUTE, T.LANE_AUX_DIM), |
| "polygon_x": (1, C.NUM_POLYGONS, T.MIXER_T["polygon"], T.MIXER_CIN["polygon"]), |
| "line_string_x": (1, C.NUM_LINE_STRINGS, T.MIXER_T["line_string"], T.MIXER_CIN["line_string"]), |
| "goal_x": (1, 1, 1, C.POSE_DIM), |
| "ego_shape_x": (1, 1, 1, C.EGO_SHAPE_DIM), |
| "turn_x": (1, 1, 1, C.TURN_INDICATOR_HISTORY), |
| "token_valid": (1, 1, T.TOKENS, 1), |
| "pos_aug": (1, 1, T.TOKENS, T.POS_AUG_DIM), |
| "fusion_key_row": (1, 1, 1, T.TOKENS), |
| "agent_key_row": (1, 1, 1, T.AGENTS), |
| "cs": (1, 1, T.AGENTS, CS_COLS), |
| "y0": (1, 1, T.AGENTS, T.STATE_COLS), |
| } |
| BF16_INPUTS = ("fusion_key_row", "agent_key_row") |
|
|
|
|
| def decoder_state(x: np.ndarray, current_states: Optional[np.ndarray] = None) -> np.ndarray: |
| """A ``[321, 81, 4]`` state -> ``[1, 1, 352, 324]`` float32 (extra rows zero). With ``current_states`` the t = 0 |
| columns are replaced by them (the ``cs`` input); with None they are zeroed (the ``y`` state).""" |
| flat = np.asarray(x, f32).reshape(C.MAX_NUM_AGENTS, T.STATE_COLS) |
| out = pad_rows(flat, T.AGENTS).astype(f32) |
| out[:, :T.STATE_COLS_T0] = 0.0 |
| if current_states is not None: |
| out[:C.MAX_NUM_AGENTS, :T.STATE_COLS_T0] = np.asarray(current_states, f32) |
| return out[None, None] |
|
|
|
|
| def plan_inputs(prep: Prepared) -> Dict[str, np.ndarray]: |
| """``{input name: array}`` for one plan (see :data:`INPUT_SPECS`).""" |
| f = prep.features |
| ego_rows, nb_rows = list(T.ISLAND_ROWS["ego"]), list(T.ISLAND_ROWS["neighbor"]) |
| tv = pad_rows(np.asarray(f.token_valid, bool), T.TOKENS) |
| nb_type = np.asarray(f.neighbor_type, f32) |
| out = { |
| "ego_x": np.asarray(f.ego, f32)[ego_rows][None, None], |
| "neighbor_x": np.asarray(f.neighbor, f32)[:, nb_rows][None], |
| "neighbor_aux": np.concatenate([nb_type, np.ones((nb_type.shape[0], 1), f32)], 1)[None, None], |
| "static_x": np.asarray(f.static, f32)[None, None], |
| "lane_x": np.asarray(f.lane, f32)[None], |
| "lane_aux": lane_aux_features(f.lane_speed, f.lane_has_speed, f.lane_attr)[None, None], |
| "route_x": np.asarray(f.route, f32)[None], |
| "route_aux": lane_aux_features(f.route_speed, f.route_has_speed, f.route_attr)[None, None], |
| "polygon_x": np.asarray(f.polygon, f32)[None], |
| "line_string_x": np.asarray(f.line_string, f32)[None], |
| "goal_x": np.asarray(f.goal, f32).reshape(1, 1, 1, -1), |
| "ego_shape_x": np.asarray(f.ego_shape, f32).reshape(1, 1, 1, -1), |
| "turn_x": np.asarray(f.turn, f32).reshape(1, 1, 1, -1), |
| "token_valid": tv.astype(f32).reshape(1, 1, -1, 1), |
| "pos_aug": np.concatenate([pad_rows(np.asarray(f.pos, f32), T.TOKENS), tv.astype(f32)[:, None]], |
| 1)[None, None], |
| "fusion_key_row": key_bias_row(pad_rows(np.asarray(f.key_valid, bool), T.TOKENS)), |
| "agent_key_row": key_bias_row(pad_rows(np.asarray(prep.decoder.agent_valid, bool), T.AGENTS)), |
| "cs": decoder_state(np.zeros((C.MAX_NUM_AGENTS, C.OUTPUT_T + 1, C.POSE_DIM), f32), |
| prep.decoder.current_states)[..., :CS_COLS], |
| "y0": decoder_state(prep.x_T), |
| } |
| for name, arr in out.items(): |
| if arr.shape != INPUT_SPECS[name]: |
| raise AssertionError(f"input {name}: shape {arr.shape}, expected {INPUT_SPECS[name]}") |
| return out |
|
|
|
|
| def warmup_inputs() -> Dict[str, np.ndarray]: |
| """Valid-looking initial contents (every token / agent a valid key: no fully masked softmax row).""" |
| out = {name: np.zeros(shape, f32) for name, shape in INPUT_SPECS.items()} |
| out["fusion_key_row"] = key_bias_row(np.arange(T.TOKENS) < T.TOKENS_REAL) |
| out["agent_key_row"] = key_bias_row(np.arange(T.AGENTS) < T.AGENTS_REAL) |
| out["token_valid"][..., :T.TOKENS_REAL, 0] = 1.0 |
| return out |
|
|
