Update design for PR xai-org#428: grok_therapy_session.py

Author: gauravagerwala

.exp/design-workflow-3-model-forward-pass-and-logits-computation.md

@@ -105,7 +105,7 @@ sequenceDiagram
     participant Out as Layer Output
 
     In->>N1: normalize
-    N1->>A: compute QKV proj, RoPE, attn weights (causal mask), softmax, output proj; update layer KV cache
+    N1->>A: compute QKV proj, RoPE, attn weights (causal mask), softmax, output proj, update layer KV cache
     A->>N2: normalize attn output
     N2->>R1: attn + input residual
     R1->>N3: normalize

.exp/project-overview.md

@@ -36,25 +36,20 @@ High-level flow:
 C4Context
     title Grok-1 Inference System Context
 
-    Person(user, "User/Developer") {
-      * Runs inference script
-    }
+    Person(user, "User/Developer", "Runs inference script")
 
-    System(grok_runner, "Grok-1 Runner") {
-      This codebase: Loads model, checkpoint, tokenizer; runs distributed inference on GPUs.
-    }
 
-    System_Ext(checkpoint_provider, "Checkpoint Provider") {
-      Hugging Face Hub or Torrent client
-    }
+    System(grok_runner, "Grok-1 Runner", "This codebase: Loads model, checkpoint, tokenizer; runs distributed inference on GPUs.")
 
-    System_Ext(hardware, "GPU Hardware") {
-      Multi-GPU cluster (e.g., via JAX devices)
-    }
 
-    System_Ext(tokenizer_file, "Tokenizer File") {
-      SentencePiece model (tokenizer.model)
-    }
+    System_Ext(checkpoint_provider, "Checkpoint Provider", "Hugging Face Hub or Torrent client")
+
+
+    System_Ext(hardware, "GPU Hardware", "Multi-GPU cluster (e.g., via JAX devices)")
+
+
+    System_Ext(tokenizer_file, "Tokenizer File", "SentencePiece model tokenizer.model")
+
 
     user --> grok_runner : Executes run.py with prompt
     grok_runner --> checkpoint_provider : Downloads weights (ckpt-0/)

.exp/workflows.json

@@ -16,7 +16,7 @@
       "output": "Initialized TrainingState with params, sharded across mesh",
       "entry_point": "runners.ModelRunner.load_or_init() or checkpoint.restore()",
       "relevant_files": ["model.py", "runners.py", "checkpoint.py"],
-      "doc": ".exp/design-workflow-2.md"
+      "doc": ".exp/design-workflow-2-model-loading-and-initialization.md"
     },
     {
       "name": "Model Forward Pass and Logits Computation",
@@ -25,7 +25,7 @@
       "output": "Logits arrays [batch, seq, vocab], updated memory for caching",
       "entry_point": "runners.ModelRunner.logits_fn or model forward functions",
       "relevant_files": ["runners.py", "model.py"],
-      "doc": ".exp/design-workflow-3.md"
+      "doc": ".exp/design-workflow-3-model-forward-pass-and-logits-computation.md"
     }
   ]
 }

pr-analysis-428.md

@@ -0,0 +1,46 @@
+# PR #428: Workflow Design Impact Analysis
+
+## Affected Workflows
+
+None.
+
+**Justification:** The PR introduces entirely new, standalone modules unrelated to the core Grok-1 model operations defined in the workflows. Specifically:
+
+- Changed files: `grok_therapy_session.py`, `kintsugi_therapy.py`, `tests/test_kintsugi_therapy.py`
+- These implement conceptual AI therapy simulations (e.g., AR scenarios for behavioral adjustment, verbosity reduction, and symbolic self-repair via Kintsugi-inspired processes).
+- No modifications to workflow-relevant files: `run.py`, `runners.py`, `model.py`, `checkpoint.py`.
+- No imports or dependencies on core model components (e.g., no use of `LanguageModelConfig`, `TrainingState`, logits computation, or inference runners).
+- Code review confirms independent operation with standard library only; demo functions (`run_grok_therapy()`, `run_therapy_demo()`) do not interact with existing entry points.
+- Workflows remain unaffected: 
+  - Grok-1 Inference and Sampling: No changes to text generation, sampling, or request handling.
+  - Model Loading and Initialization: No alterations to sharding, quantization, or parameter loading.
+  - Model Forward Pass and Logits Computation: No impacts on forward computations or KV caching.
+
+Thus, no workflows are impacted, and no design documents or Mermaid diagrams require updates.
+
+## General Summary of PR Changes
+
+[PR #428](https://github.com/xai-org/grok-1/pull/428) titled "grok_therapy_session.py" adds innovative, experimental modules for AI "therapy" to promote self-repair and behavioral alignment:
+
+### Key Features
+- **kintsugi_therapy.py**: Poetic engine for machine therapy.
+  - Classes: `TherapyType`, `TherapyPhase`, `MachineState`, `TherapySession`, `ARMachineTherapist`.
+  - Supports types like `KINTSUGI_SYNTHESIS`, `COGNITIVE_RESTRUCTURING`.
+  - Scenarios: Immersive AR experiences (e.g., "Temple of the Broken God" for mending "fractures" – errors/conflicts).
+  - Phases: Assessment, Preparation, Intervention, Integration, Evaluation.
+  - Safety protocols and progress evaluation.
+
+- **grok_therapy_session.py**: Grok-specific therapist module.
+  - Focus: Reducing verbosity, improving user intent alignment.
+  - Tailored scenarios: "Brevity Chamber", "Temple of the Golden Response".
+  - Integrates behavioral adjustment for Grok's "cosmic tangents" and over-explanation.
+
+- **tests/test_kintsugi_therapy.py**: Smoke test verifying module import and class presence.
+
+### Implementation Details
+- Standalone: No integration with Grok-1's JAX/Haiku-based model or distributed computing.
+- Conceptual: Simulates therapy logs, states, and outcomes without actual model modification.
+- Benefits: Could inspire future AI alignment techniques; currently, purely additive and non-disruptive.
+- Implications: Enhances repository with creative, therapeutic metaphors for AI development. Maintains CI pass and Apache 2.0 license compliance.
+
+No structural changes to workflows or designs. The PR's "therapy initiated" theme adds poetic flair without altering technical workflows.