Fix review findings: consistent naming, public imports, lint, docstring

- Rename graph_builder to builder in _gemma4.py and
  _vision_language_3model.py for consistency with other 22 task files
- Change GraphBuilder import from onnxscript._internal.builder (private)
  to onnxscript (public API) in _base.py, _cache_utils.py, _causal_lm.py,
  _gemma4.py
- Fix import sorting and line length issues flagged by lintrunner
- Update _make_graph() docstring to mention builder.add_output()

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>

Author: justinchuby

src/mobius/tasks/_adapter.py

@@ -29,10 +29,18 @@ def build(
         # Determine input shape based on adapter type
         if hasattr(config, "in_channels"):
             # T2I-Adapter: conditioning image input
-            condition = builder.input("condition", dtype=ir.DataType.FLOAT, shape=["batch", config.in_channels, "height", "width"])
+            condition = builder.input(
+                "condition",
+                dtype=ir.DataType.FLOAT,
+                shape=["batch", config.in_channels, "height", "width"],
+            )
         else:
             # IP-Adapter: image embedding input
-            condition = builder.input("image_embeds", dtype=ir.DataType.FLOAT, shape=["batch", config.image_embed_dim])
+            condition = builder.input(
+                "image_embeds",
+                dtype=ir.DataType.FLOAT,
+                shape=["batch", config.image_embed_dim],
+            )
 
         outputs = module(op, condition)
 

src/mobius/tasks/_audio_feature_extraction.py

@@ -31,7 +31,9 @@ def build(
         graph, builder = _make_graph()
         op = builder.op
 
-        input_values = builder.input("input_values", dtype=ir.DataType.FLOAT, shape=["batch", "time"])
+        input_values = builder.input(
+            "input_values", dtype=ir.DataType.FLOAT, shape=["batch", "time"]
+        )
 
         last_hidden_state = module(op, input_values=input_values)
 

src/mobius/tasks/_base.py

@@ -9,8 +9,7 @@
 from typing import ClassVar
 
 import onnx_ir as ir
-from onnxscript import nn
-from onnxscript._internal.builder import GraphBuilder
+from onnxscript import GraphBuilder, nn
 
 import mobius
 from mobius._configs import BaseModelConfig
@@ -108,6 +107,7 @@ def _make_graph(
     """Create an empty graph and its builder.
 
     Inputs should be added after creation via ``builder.input()``.
+    Outputs should be registered via ``builder.add_output()``.
 
     Returns:
         ``(graph, builder)`` — call ``builder.op`` to get the op handle.

src/mobius/tasks/_cache_utils.py

@@ -14,7 +14,7 @@
 from typing import NamedTuple
 
 import onnx_ir as ir
-from onnxscript._internal.builder import GraphBuilder
+from onnxscript import GraphBuilder
 
 from mobius._configs import BaseModelConfig
 

src/mobius/tasks/_causal_lm.py

@@ -6,8 +6,7 @@
 from __future__ import annotations
 
 import onnx_ir as ir
-from onnxscript import nn
-from onnxscript._internal.builder import GraphBuilder
+from onnxscript import GraphBuilder, nn
 
 from mobius._configs import ArchitectureConfig
 from mobius._model_package import ModelPackage
@@ -121,7 +120,9 @@ def build(
         # --- Cache setup (static vs dynamic) ---
         if static:
             attention_mask = None
-            position_ids = builder.input("position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len])
+            position_ids = builder.input(
+                "position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len]
+            )
             past_key_values = _make_static_cache_inputs(
                 builder,
                 config.num_hidden_layers,
@@ -134,9 +135,13 @@ def build(
         else:
             past_seq_len = ir.SymbolicDim("past_sequence_len")
             attention_mask = builder.input(
-                "attention_mask", dtype=ir.DataType.INT64, shape=[batch, "past_seq_len + seq_len"],
+                "attention_mask",
+                dtype=ir.DataType.INT64,
+                shape=[batch, "past_seq_len + seq_len"],
+            )
+            position_ids = builder.input(
+                "position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len]
             )
-            position_ids = builder.input("position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len])
 
             # MLA attention: K/V heads equal q heads (no GQA reduction in
             # latent space).  The ONNX Attention op is called with
@@ -222,9 +227,13 @@ def build(
 
         input_ids = builder.input("input_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len])
         attention_mask = builder.input(
-            "attention_mask", dtype=ir.DataType.INT64, shape=[batch, "past_seq_len + seq_len"],
+            "attention_mask",
+            dtype=ir.DataType.INT64,
+            shape=[batch, "past_seq_len + seq_len"],
+        )
+        position_ids = builder.input(
+            "position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len]
         )
-        position_ids = builder.input("position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len])
 
         past_key_values = _make_hybrid_cache_inputs(
             builder,
@@ -289,10 +298,14 @@ def _make_static_cache_inputs(
 
     # Shared inputs across all layers
     write_indices = builder.input(
-        "write_indices", dtype=ir.DataType.INT64, shape=[batch],
+        "write_indices",
+        dtype=ir.DataType.INT64,
+        shape=[batch],
     )
     nonpad_kv_seqlen = builder.input(
-        "nonpad_kv_seqlen", dtype=ir.DataType.INT64, shape=[batch],
+        "nonpad_kv_seqlen",
+        dtype=ir.DataType.INT64,
+        shape=[batch],
     )
 
     # Build StaticCacheState for each layer (shared indices)

src/mobius/tasks/_codec.py

@@ -89,7 +89,9 @@ def _build_encoder(
         audio_len = ir.SymbolicDim("audio_length")
 
         graph, builder = _make_graph(name="encoder")
-        waveform = builder.input("waveform", dtype=ir.DataType.FLOAT, shape=[batch, 1, audio_len])
+        waveform = builder.input(
+            "waveform", dtype=ir.DataType.FLOAT, shape=[batch, 1, audio_len]
+        )
 
         codes = encoder(builder.op, waveform)
 

src/mobius/tasks/_controlnet.py

@@ -31,10 +31,22 @@ def build(
         graph, builder = _make_graph()
         op = builder.op
 
-        sample = builder.input("sample", dtype=ir.DataType.FLOAT, shape=["batch", config.in_channels, "height", "width"])
+        sample = builder.input(
+            "sample",
+            dtype=ir.DataType.FLOAT,
+            shape=["batch", config.in_channels, "height", "width"],
+        )
         timestep = builder.input("timestep", dtype=ir.DataType.INT64, shape=["batch"])
-        encoder_hidden_states = builder.input("encoder_hidden_states", dtype=ir.DataType.FLOAT, shape=["batch", "sequence_length", config.cross_attention_dim])
-        controlnet_cond = builder.input("controlnet_cond", dtype=ir.DataType.FLOAT, shape=["batch", config.conditioning_channels, "cond_height", "cond_width"])
+        encoder_hidden_states = builder.input(
+            "encoder_hidden_states",
+            dtype=ir.DataType.FLOAT,
+            shape=["batch", "sequence_length", config.cross_attention_dim],
+        )
+        controlnet_cond = builder.input(
+            "controlnet_cond",
+            dtype=ir.DataType.FLOAT,
+            shape=["batch", config.conditioning_channels, "cond_height", "cond_width"],
+        )
 
         down_outputs, mid_output = module(
             op,

src/mobius/tasks/_denoising.py

@@ -31,9 +31,17 @@ def build(
         graph, builder = _make_graph()
         op = builder.op
 
-        sample = builder.input("sample", dtype=ir.DataType.FLOAT, shape=["batch", config.in_channels, "height", "width"])
+        sample = builder.input(
+            "sample",
+            dtype=ir.DataType.FLOAT,
+            shape=["batch", config.in_channels, "height", "width"],
+        )
         timestep = builder.input("timestep", dtype=ir.DataType.INT64, shape=["batch"])
-        encoder_hidden_states = builder.input("encoder_hidden_states", dtype=ir.DataType.FLOAT, shape=["batch", "sequence_length", config.cross_attention_dim])
+        encoder_hidden_states = builder.input(
+            "encoder_hidden_states",
+            dtype=ir.DataType.FLOAT,
+            shape=["batch", "sequence_length", config.cross_attention_dim],
+        )
 
         noise_pred = module(
             op,

src/mobius/tasks/_feature_extraction.py

@@ -41,8 +41,12 @@ def build(
         op = builder.op
 
         input_ids = builder.input("input_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len])
-        attention_mask = builder.input("attention_mask", dtype=ir.DataType.INT64, shape=[batch, seq_len])
-        token_type_ids = builder.input("token_type_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len])
+        attention_mask = builder.input(
+            "attention_mask", dtype=ir.DataType.INT64, shape=[batch, seq_len]
+        )
+        token_type_ids = builder.input(
+            "token_type_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len]
+        )
 
         last_hidden_state = module(
             op,