{
  "nbformat": 4,
  "nbformat_minor": 5,
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3"
    },
    "article_scope": "This teaching model isolates the three matching cases in the figure. The production loader also validates ordered entries, handles unknown or misplaced stages, and works with the runtime checker. Its linked source and tests are the authority."
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Repeated stage matching\n",
        "This teaching model isolates the three matching cases in the figure. The production loader also validates ordered entries, handles unknown or misplaced stages, and works with the runtime checker. Its linked source and tests are the authority.\n",
        "\nSources:\n",
        "- https://github.com/torch-spyre/aiu-trace-analyzer/pull/98/files\n",
        "- https://github.com/torch-spyre/aiu-trace-analyzer/blob/0b4a5e538fc19e03cb9fdc6c6bed57a04e03ed9f/src/aiu_trace_analyzer/core/stage_profile.py\n"
      ],
      "id": "cell-1"
    },
    {
      "cell_type": "code",
      "metadata": {},
      "execution_count": 1,
      "source": [
        "# A small teaching model of the three displayed cases, not the full loader.\n",
        "def repeated_stage_flags(occurrences, key=None, complete=None):\n",
        "    if complete is not None:\n",
        "        if len(complete) != occurrences:\n",
        "            raise ValueError(\"Complete profile must supply every displayed occurrence\")\n",
        "        return list(complete)\n",
        "    if key == \"pipeline_barrier\":\n",
        "        return [True] * occurrences\n",
        "    prefix = \"pipeline_barrier#\"\n",
        "    if not isinstance(key, str) or not key.startswith(prefix):\n",
        "        raise ValueError(\"Use pipeline_barrier or pipeline_barrier#N\")\n",
        "    ordinal = int(key[len(prefix):])\n",
        "    if not 1 <= ordinal <= occurrences:\n",
        "        raise ValueError(\"Occurrence is out of range\")\n",
        "    return [i == ordinal for i in range(1, occurrences + 1)]\n",
        "\n",
        "assert repeated_stage_flags(4, \"pipeline_barrier\") == [True] * 4\n",
        "assert repeated_stage_flags(4, \"pipeline_barrier#2\") == [False, True, False, False]\n",
        "assert repeated_stage_flags(4, complete=[False, True, True, True]) == [False, True, True, True]\n",
        "for label, flags in [\n",
        "    (\"bare\", repeated_stage_flags(4, \"pipeline_barrier\")),\n",
        "    (\"#2\", repeated_stage_flags(4, \"pipeline_barrier#2\")),\n",
        "    (\"complete\", repeated_stage_flags(4, complete=[False, True, True, True])),\n",
        "]:\n",
        "    print(label, flags)\n"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "bare [True, True, True, True]\n",
            "#2 [False, True, False, False]\n",
            "complete [False, True, True, True]\n"
          ]
        }
      ],
      "id": "cell-2"
    }
  ]
}