weco 0.2.5__py3-none-any.whl → 0.2.6__py3-none-any.whl

weco/__init__.py

@@ -1,4 +1,4 @@
 # DO NOT EDIT
-__pkg_version__ = "0.2.5"
+__pkg_version__ = "0.2.6"
 __api_version__ = "v1"
 __base_url__ = f"https://api.aide.weco.ai/{__api_version__}"

weco/cli.py

@@ -321,7 +321,12 @@ def main() -> None:
             _, best_solution_panel = solution_panels.get_display(current_step=steps)
 
             # Update the end optimization layout
-            end_optimization_layout["summary"].update(summary_panel.get_display())
+            final_message = (
+                f"{summary_panel.metric_name.capitalize()} {'maximized' if summary_panel.maximize else 'minimized'}! Best solution {summary_panel.metric_name.lower()} = [green]{status_response['best_result']['metric_value']}[/] 🏆"
+                if best_solution_node is not None
+                else "[red] No solution found.[/]"
+            )
+            end_optimization_layout["summary"].update(summary_panel.get_display(final_message=final_message))
             end_optimization_layout["tree"].update(tree_panel.get_display())
             end_optimization_layout["best_solution"].update(best_solution_panel)
 

weco/panels.py

@@ -12,7 +12,9 @@ class SummaryPanel:
     """Holds a summary of the optimization session."""
 
     def __init__(self, maximize: bool, metric_name: str, total_steps: int, model: str, session_id: str = None):
-        self.goal = ("Maximizing" if maximize else "Minimizing") + f" {metric_name}..."
+        self.maximize = maximize
+        self.metric_name = metric_name
+        self.goal = ("Maximizing" if self.maximize else "Minimizing") + f" {self.metric_name}..."
         self.total_input_tokens = 0
         self.total_output_tokens = 0
         self.total_steps = total_steps
@@ -39,24 +41,28 @@ class SummaryPanel:
         self.total_input_tokens += usage["input_tokens"]
         self.total_output_tokens += usage["output_tokens"]
 
-    def get_display(self) -> Panel:
+    def get_display(self, final_message: Optional[str] = None) -> Panel:
         """Create a summary panel with the relevant information."""
         layout = Layout(name="summary")
         summary_table = Table(show_header=False, box=None, padding=(0, 1))
         # Goal
-        summary_table.add_row(f"[bold cyan]Goal:[/] {self.goal}")
+        if final_message is not None:
+            summary_table.add_row(f"[bold cyan]Result:[/] {final_message}")
+        else:
+            summary_table.add_row(f"[bold cyan]Goal:[/] {self.goal}")
+        summary_table.add_row("")
+        # Model used
+        summary_table.add_row(f"[bold cyan]Model:[/] {self.model}")
         summary_table.add_row("")
         # Log directory
         runs_dir = f".runs/{self.session_id}"
         summary_table.add_row(f"[bold cyan]Logs:[/] [blue underline]{runs_dir}[/]")
         summary_table.add_row("")
-        # Model used
-        summary_table.add_row(f"[bold cyan]Model:[/] [yellow]{self.model}[/]")
-        summary_table.add_row("")
         # Token counts
         summary_table.add_row(
             f"[bold cyan]Tokens:[/] ↑[yellow]{format_number(self.total_input_tokens)}[/] ↓[yellow]{format_number(self.total_output_tokens)}[/] = [green]{format_number(self.total_input_tokens + self.total_output_tokens)}[/]"
         )
+        summary_table.add_row("")
         # Progress bar
         summary_table.add_row(self.progress)
 

{weco-0.2.5.dist-info → weco-0.2.6.dist-info}/METADATA

@@ -1,8 +1,8 @@
 Metadata-Version: 2.4
 Name: weco
-Version: 0.2.5
+Version: 0.2.6
 Summary: Documentation for `weco`, a CLI for using Weco AI's code optimizer.
-Author-email: Weco AI Team <dhruv@weco.ai>
+Author-email: Weco AI Team <contact@weco.ai>
 License: MIT
 Project-URL: Homepage, https://github.com/WecoAI/weco-cli
 Keywords: AI,Code Optimization,Code Generation
@@ -99,32 +99,69 @@ Here's how `weco` can be applied to common ML engineering tasks:
 
 ### Examples
 
-**Example 1: Optimizing PyTorch operations**
+**Example 1: Optimizing PyTorch simple operations**
 
 ```bash
-weco --source examples/simple-torch/optimize.py \
-     --eval-command "python examples/simple-torch/evaluate.py --solution-path examples/simple-torch/optimize.py --device mps" \
+cd examples/hello-kernel-world
+pip install torch 
+weco --source optimize.py \
+     --eval-command "python evaluate.py --solution-path optimize.py --device cpu" \
      --metric speedup \
      --maximize true \
      --steps 15 \
-     --model o3-mini \
+     --model claude-3-7-sonnet-20250219 \
      --additional-instructions "Fuse operations in the forward method while ensuring the max float deviation remains small. Maintain the same format of the code."
 ```
 
+Note that if you have an NVIDIA gpu, change the device to `cuda`. If you are running this on Apple Silicon, set it to `mps`.
+
 **Example 2: Optimizing MLX operations with instructions from a file**
 
-Sometimes, additional context or instructions are too complex for a single command-line string. You can provide a path to a file containing these instructions.
+Lets optimize a 2D convolution operation in [`mlx`](https://github.com/ml-explore/mlx) using [Metal](https://developer.apple.com/documentation/metal/). Sometimes, additional context or instructions are too complex for a single command-line string. You can provide a path to a file containing these instructions.
 
 ```bash
-weco --source examples/simple-mlx/optimize.py \
-     --eval-command "python examples/simple-mlx/evaluate.py --solution-path examples/simple-mlx/optimize.py" \
+cd examples/metal
+pip install mlx
+weco --source optimize.py \
+     --eval-command "python evaluate.py --solution-path optimize.py" \
      --metric speedup \
      --maximize true \
      --steps 30 \
      --model o3-mini \
-     --additional-instructions examples/simple-mlx/metal-examples.rst
+     --additional-instructions examples.rst
 ```
 
+**Example 3: Level Agnostic Optimization: Causal Self Attention with Triton & CUDA**
+
+Given how useful causal multihead self attention is to transformers, we've seen its wide adoption across ML engineering and AI research. Its great to keep things at a high-level (in PyTorch) when doing research, but when moving to production you often need to write highly customized low-level kernels to make things run as fast as they can. The `weco` CLI can optimize kernels across a variety of different abstraction levels and frameworks. Example 2 uses Metal but lets explore two more frameworks:
+
+1. [Triton](https://github.com/triton-lang/triton)
+    ```bash
+   cd examples/triton
+   pip install torch triton
+   weco --source optimize.py \
+        --eval-command "python evaluate.py --solution-path optimize.py" \
+        --metric speedup \
+        --maximize true \
+        --steps 30 \
+        --model gemini-2.5-pro-preview-03-25 \
+        --additional-instructions "Use triton to optimize the code while ensuring a small max float diff. Maintain the same code format."
+   ```
+
+2. [CUDA](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html)
+   ```bash
+   cd examples/cuda
+   pip install torch
+   weco --source optimize.py \
+        --eval-command "python evaluate.py --solution-path optimize.py" \
+        --metric speedup \
+        --maximize true \
+        --steps 30 \
+        --model gemini-2.5-pro-preview-03-25 \
+        --additional-instructions guide.md
+   ```
+
+
 ---
 
 ### Command Line Arguments

weco-0.2.6.dist-info/RECORD

@@ -0,0 +1,11 @@
+weco/__init__.py,sha256=a3RxrwZhsuinalG_NtT0maKLFXFbgmgXFahpFgcEtZQ,124
+weco/api.py,sha256=8rIf2Fy3tN6GW7BG1CaggtfE9pW56I1erzwLCgawcVE,3511
+weco/cli.py,sha256=6rGEm_L-WSkJIT-jgfFmf2i_DXkQn6ILhqYQlptLFew,17159
+weco/panels.py,sha256=9gq5C43hgUmQgl6tW-f2dBbDjlsBKBatSaUVKeGm4Zw,12296
+weco/utils.py,sha256=hhIebUPnetFMfNSFfcsKVw1TSpeu_Zw3rBPPnxDie0U,3911
+weco-0.2.6.dist-info/licenses/LICENSE,sha256=p_GQqJBvuZgkLNboYKyH-5dhpTDlKs2wq2TVM55WrWE,1065
+weco-0.2.6.dist-info/METADATA,sha256=Ih-nkHxq_SJKccYXoVHmxytElqG75BSU1DGAAC9ipkk,11581
+weco-0.2.6.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+weco-0.2.6.dist-info/entry_points.txt,sha256=ixJ2uClALbCpBvnIR6BXMNck8SHAab8eVkM9pIUowcs,39
+weco-0.2.6.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
+weco-0.2.6.dist-info/RECORD,,

weco-0.2.5.dist-info/RECORD

@@ -1,11 +0,0 @@
-weco/__init__.py,sha256=uqfsknG6Jxky4ddpuNn9SYN2lWNJloPDykqCjcmU1UQ,124
-weco/api.py,sha256=8rIf2Fy3tN6GW7BG1CaggtfE9pW56I1erzwLCgawcVE,3511
-weco/cli.py,sha256=rb01pl0Q8sLKBpZ9Qtlz044pEppJAa8CJkv2XM9zzFo,16753
-weco/panels.py,sha256=5HrDrYzx_vDjV41mlZmOOZYxdJIXD5RzBBNKXxUOLEY,12022
-weco/utils.py,sha256=hhIebUPnetFMfNSFfcsKVw1TSpeu_Zw3rBPPnxDie0U,3911
-weco-0.2.5.dist-info/licenses/LICENSE,sha256=p_GQqJBvuZgkLNboYKyH-5dhpTDlKs2wq2TVM55WrWE,1065
-weco-0.2.5.dist-info/METADATA,sha256=bjW2DWItTu3y_lmLllZIRz43hh5vMCi_u4FXW5eKz5E,9863
-weco-0.2.5.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-weco-0.2.5.dist-info/entry_points.txt,sha256=ixJ2uClALbCpBvnIR6BXMNck8SHAab8eVkM9pIUowcs,39
-weco-0.2.5.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
-weco-0.2.5.dist-info/RECORD,,