Fix Visualizer tab bugs: softmax data, reset button, causal toggle

app.py

@@ -159,9 +159,13 @@ def create_app() -> gr.Blocks:
                     """Move to previous step."""
                     return max(current_step - 1, 0)
                 
-                def reset_step():
-                    """Reset to step 0."""
-                    return 0
+                def reset_step(seq_len, block_size, causal):
+                    """Reset to step 0 and update visualizations."""
+                    step = 0
+                    tiling_fig = create_tiling_grid(seq_len, block_size, step, causal)
+                    num_tiles = seq_len // block_size
+                    softmax_fig, explanation = create_online_softmax_state(step, num_tiles)
+                    return step, tiling_fig, softmax_fig, explanation
                 
                 # Wire up events
                 viz_inputs = [seq_len_viz, block_size_viz, causal_viz, step_slider]
@@ -170,7 +174,12 @@ def create_app() -> gr.Blocks:
                 # Update on parameter change
                 seq_len_viz.change(fn=update_visualizations, inputs=viz_inputs, outputs=viz_outputs)
                 block_size_viz.change(fn=update_visualizations, inputs=viz_inputs, outputs=viz_outputs)
-                causal_viz.change(fn=update_visualizations, inputs=viz_inputs, outputs=viz_outputs)
+                # Causal toggle resets step to 0 since the grid structure changes
+                causal_viz.change(
+                    fn=reset_step,
+                    inputs=[seq_len_viz, block_size_viz, causal_viz],
+                    outputs=[step_slider, tiling_plot, softmax_plot, softmax_explanation]
+                )
                 step_slider.change(fn=update_visualizations, inputs=viz_inputs, outputs=viz_outputs)
                 
                 # Step controls
@@ -180,7 +189,11 @@ def create_app() -> gr.Blocks:
                     outputs=step_slider
                 )
                 step_back_btn.click(fn=step_back, inputs=step_slider, outputs=step_slider)
-                reset_btn.click(fn=reset_step, outputs=step_slider)
+                reset_btn.click(
+                    fn=reset_step,
+                    inputs=[seq_len_viz, block_size_viz, causal_viz],
+                    outputs=[step_slider, tiling_plot, softmax_plot, softmax_explanation]
+                )
                 
                 # Memory hierarchy
                 algo_choice.change(fn=update_memory_hierarchy, inputs=algo_choice, outputs=memory_plot)

src/visualizer.py

@@ -194,7 +194,7 @@ def create_online_softmax_state(
             "m_before": 2.1,
             "m_after": 3.5,
             "l_before": 3.42,
-            "l_after": 1.70,  # 3.42 * exp(2.1-3.5) + 5.21 = 0.85 + 5.21
+            "l_after": 6.06,  # 3.42 * exp(2.1-3.5) + 5.21 = 0.85 + 5.21 ≈ 6.06
             "rescale_factor": 0.247,  # exp(2.1 - 3.5)
             "rescaled": True,
         },