Replace op.Constant(value_float=...) with Python float literals, use autocast (#166)

Continuation of PR #58 which established this pattern.

Replace all 68 instances of op.Constant(value_float=X) across 30 files
with plain Python float literals or float() expressions. This leverages
onnxscript's auto-casting of Python scalars to match the dtype of the
other operand in binary ops, which fixes dtype mismatches when models
use bfloat16 or float16.

Patterns replaced:
- op.Constant(value_float=X) -> X (plain literal)
- op.Constant(value_float=float(expr)) -> float(expr)
- op.Constant(value_float=self.attr) -> self.attr
- op.CastLike(op.Constant(value_float=X), ref) -> op.CastLike(X, ref)

Note: This is part 1. (A similar change can be done for int/ints/floats,
but will do that in separate PRs.)

---------

Signed-off-by: G Ramalingam <grama@microsoft.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: gramalingam

src/mobius/components/_ecapa_tdnn.py

@@ -313,7 +313,7 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value):
             op.ReduceSum(op.Mul(diff, diff), [2], keepdims=True),
             seq_length_float,
         )
-        eps_const = op.Constant(value_float=self._eps)
+        eps_const = self._eps
         std = op.Sqrt(op.Add(variance, eps_const))
 
         # Expand mean and std to (batch, channels, time) for concatenation

src/mobius/components/_multimodal.py

@@ -182,7 +182,7 @@ def forward(
             axis=0,
         )
         zero_pad = op.Expand(
-            op.CastLike(op.Constant(value_float=0.0), vision_embeddings),
+            op.CastLike(0.0, vision_embeddings),
             pad_shape,
         )
         # [batch, vision_seq + 1, hidden]

src/mobius/components/_qwen25_vl_vision.py

@@ -252,7 +252,7 @@ def _emit_standard_attention(self, op, q, k, v, cu_seqlens, seq_len_val):
         v = op.Unsqueeze(v, [0])
 
         # Build block-diagonal attention bias from cu_seqlens
-        attn_bias = self._build_block_diagonal_bias(op, cu_seqlens, seq_len_val)
+        attn_bias = self._build_block_diagonal_bias(op, cu_seqlens, seq_len_val, q)
         attn_bias = op.Unsqueeze(attn_bias, [0, 1])  # (1, 1, N, N)
 
         # Scaled dot-product attention
@@ -301,11 +301,14 @@ def _apply_rotary(self, op, x, cos, sin):
 
         return op.Concat(rot_x1, rot_x2, axis=-1)
 
-    def _build_block_diagonal_bias(self, op, cu_seqlens, seq_len):
+    def _build_block_diagonal_bias(self, op, cu_seqlens, seq_len, dtype_ref):
         """Build block-diagonal attention bias from cu_seqlens.
 
         Creates a matrix where positions in the same sub-sequence have 0
         and positions in different sub-sequences have -inf.
+
+        Args:
+            dtype_ref: reference tensor whose dtype the bias should match.
         """
         # Create range indices
         indices = op.Range(
@@ -337,8 +340,9 @@ def _build_block_diagonal_bias(self, op, cu_seqlens, seq_len):
         same_segment = op.Equal(seg_row, seg_col)
 
         # Convert to attention bias: 0 where same segment, -inf where different
-        neg_inf = op.Constant(value_float=-1e9)
-        zero = op.Constant(value_float=0.0)
+        # CastLike before Where so only the scalar is cast (cheaper than post-broadcast)
+        neg_inf = op.CastLike(-1e9, dtype_ref)
+        zero = op.CastLike(0.0, dtype_ref)
         return op.Where(same_segment, zero, neg_inf)
 
 

src/mobius/components/_qwen3_vl_vision.py

@@ -292,8 +292,8 @@ def _emit_standard_attention(self, op, query, key, value, cu_seqlens, hidden_sta
         same_segment = op.Equal(segment_row, segment_column)  # (total_seq, total_seq)
         attn_bias = op.Where(
             same_segment,
-            op.Constant(value_float=0.0),
-            op.Constant(value_float=-10000.0),
+            op.CastLike(0.0, query),
+            op.CastLike(-10000.0, query),
         )
         # Reshape for Attention: (1, 1, total_seq, total_seq)
         attn_bias = op.Unsqueeze(attn_bias, [0, 1])
@@ -609,12 +609,12 @@ def _interpolate_pos_embed(self, op, grid_thw):
         )
 
         h_idxs = body_op.Div(
-            body_op.Mul(h_range, body_op.Constant(value_float=n_minus_1)),
-            body_op.Sub(H_f, body_op.Constant(value_float=1.0)),
+            body_op.Mul(h_range, n_minus_1),
+            body_op.Sub(H_f, 1.0),
         )
         w_idxs = body_op.Div(
-            body_op.Mul(w_range, body_op.Constant(value_float=n_minus_1)),
-            body_op.Sub(W_f, body_op.Constant(value_float=1.0)),
+            body_op.Mul(w_range, n_minus_1),
+            body_op.Sub(W_f, 1.0),
         )
 
         # Floor/ceil indices
@@ -648,8 +648,8 @@ def _interpolate_pos_embed(self, op, grid_thw):
         idx_10 = body_op.Reshape(body_op.Add(bh_c2, wf2), [-1])
         idx_11 = body_op.Reshape(body_op.Add(bh_c2, wc2), [-1])
 
-        one_minus_dh = body_op.Sub(body_op.Constant(value_float=1.0), dh)
-        one_minus_dw = body_op.Sub(body_op.Constant(value_float=1.0), dw)
+        one_minus_dh = body_op.Sub(1.0, dh)
+        one_minus_dw = body_op.Sub(1.0, dw)
         dh2 = body_op.Unsqueeze(dh, [1])
         omdh2 = body_op.Unsqueeze(one_minus_dh, [1])
         dw2 = body_op.Unsqueeze(dw, [0])
@@ -704,7 +704,7 @@ def _interpolate_pos_embed(self, op, grid_thw):
             body_op.Constant(value_ints=[0]),
             axis=0,
         )
-        padded = body_op.Pad(pos_embeds, pads, body_op.Constant(value_float=0.0))
+        padded = body_op.Pad(pos_embeds, pads, 0.0)
         padded.name = "padded_pos_embed"
         body_graph.outputs.append(padded)
 

src/mobius/models/base.py

@@ -126,7 +126,7 @@ def forward(
                 to=ir.DataType.INT32,
             )  # [batch] INT32
             total_seq_len = op.Cast(
-                op.Gather(op.Shape(attention_mask), op.Constant(value_int=1)),
+                op.Gather(op.Shape(attention_mask), 1),
                 to=ir.DataType.INT32,
             )  # scalar INT32
 

src/mobius/models/blip2.py

@@ -166,7 +166,7 @@ def forward(self, op: builder.OpBuilder, input_ids: ir.Value, image_features: ir
 
         # Build indices to scatter image features into text positions
         mask_int = op.Cast(image_mask, to=7)
-        cumsum = op.CumSum(mask_int, op.Constant(value_int=1))
+        cumsum = op.CumSum(mask_int, 1)
         indices = op.Sub(cumsum, op.Constant(value_int=1))
         indices = op.Clip(indices, op.Constant(value_int=0))
 

src/mobius/models/clip.py

@@ -294,18 +294,20 @@ def forward(
         hidden_states = self.embeddings(op, input_ids)
 
         # Build causal attention bias (lower-triangular)
+        # CastLike ensures the bias dtype matches hidden_states (fp16/bf16/fp32)
+        # Cast the scalar *before* Expand so only a single element is cast
         seq_len = op.Shape(input_ids, start=1, end=2)
         _causal_mask = op.Trilu(
             op.Expand(
-                op.Constant(value_float=0.0),
+                op.CastLike(0.0, hidden_states),
                 op.Concat(seq_len, seq_len, axis=0),
             ),
             upper=0,
         )
         # Fill upper triangle with -inf
         neg_inf_mask = op.Trilu(
             op.Expand(
-                op.Constant(value_float=-10000.0),
+                op.CastLike(-10000.0, hidden_states),
                 op.Concat(seq_len, seq_len, axis=0),
             ),
             upper=1,

src/mobius/models/cogvideox.py

@@ -157,7 +157,7 @@ def forward(
             emb, num_outputs=6, axis=-1, _outputs=6
         )
 
-        one = op.Constant(value_float=1.0)
+        one = 1.0
 
         # Modulate video stream
         normed_h = self.norm(op, hidden_states)
@@ -197,7 +197,7 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value, temb: ir.Value
         emb = self.linear(op, emb)
         # CogVideoX shift-first order
         shift, scale = op.Split(emb, num_outputs=2, axis=-1, _outputs=2)
-        one = op.Constant(value_float=1.0)
+        one = 1.0
         hidden_states = self.norm(op, hidden_states)
         hidden_states = op.Mul(hidden_states, op.Add(one, op.Unsqueeze(scale, [1])))
         hidden_states = op.Add(hidden_states, op.Unsqueeze(shift, [1]))

src/mobius/models/cohere.py

@@ -127,7 +127,5 @@ def forward(
             # Scale logits by the model's configured logit_scale scalar
             # (HF default: 0.0625 = 1/16 for all Cohere models).
             # CastLike ensures the constant matches logits dtype (fp16/bf16/fp32).
-            logits = op.Mul(
-                logits, op.CastLike(op.Constant(value_float=float(self.logit_scale)), logits)
-            )
+            logits = op.Mul(logits, op.CastLike(float(self.logit_scale), logits))
         return logits, present_key_values

src/mobius/models/deepseek.py

@@ -96,7 +96,7 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value):
         # Normalize weights (V3 with norm_topk_prob=True)
         if self.norm_topk_prob:
             weight_sum = op.ReduceSum(routing_weights, [-1], keepdims=True)
-            eps = op.Constant(value_float=1e-20)
+            eps = 1e-20
             routing_weights = op.Div(routing_weights, op.Add(weight_sum, eps))
 
         # Apply routing scale