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

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/_gated_deltanet.py

@@ -80,7 +80,7 @@ def forward(
         # Zero bias — model has no conv bias; the function requires it.
         # CastLike ensures the bias matches the weight dtype (e.g. f16).
         conv_bias = op.Expand(
-            op.CastLike(op.Constant(value_float=0.0), self.weight),
+            op.CastLike(0.0, self.weight),
             op.Constant(value_ints=[self._channels]),
         )
         return op.CausalConvWithState(

src/mobius/components/_mamba_block.py

@@ -216,7 +216,7 @@ def forward(
         else:
             # Zero bias — the function requires a bias input.
             conv_bias = op.Expand(
-                op.CastLike(op.Constant(value_float=0.0), self.weight),
+                op.CastLike(0.0, self.weight),
                 op.Constant(value_ints=[self._channels]),
             )
         return op.CausalConvWithState(
@@ -393,7 +393,7 @@ def forward(
         # dt = softplus(dt_raw + dt_bias): (B, T, num_heads)
         dt = op.Softplus(op.Add(dt_raw_f32, dt_bias_f32))
         if self.time_step_min > 0.0:
-            dt = op.Clip(dt, op.Constant(value_float=self.time_step_min))
+            dt = op.Clip(dt, self.time_step_min)
 
         # decay = A * dt in log-space: g_t where exp(g_t) is the decay
         # A = -exp(A_log), so decay = -exp(A_log) * dt

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

@@ -337,8 +337,8 @@ 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)
+        neg_inf = -1e9
+        zero = 0.0
         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),
+            0.0,
+            -10000.0,
         )
         # 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/functions/causal_conv.py

@@ -113,8 +113,8 @@ def causal_conv_nd_with_state(
     conv_input = op.Concat(conv_state_val, input_val, axis=temporal_axis)
 
     # Step 2: Extract new carry state — last K-1 positions of conv_input.
-    total_len = op.Gather(op.Shape(conv_input), op.Constant(value_int=temporal_axis), axis=0)
-    state_start = op.Sub(total_len, op.Constant(value_int=state_width))
+    total_len = op.Gather(op.Shape(conv_input), temporal_axis, axis=0)
+    state_start = op.Sub(total_len, state_width)
     present_state = op.Slice(
         conv_input,
         op.Reshape(state_start, op.Constant(value_ints=[1])),

src/mobius/functions/linear_attention.py

@@ -189,7 +189,7 @@ def linear_attention(
     # CastLike ensures scale constant matches the input dtype.
     scaled_query = op.Mul(
         query_expanded,
-        op.CastLike(op.Constant(value_float=scale), query_expanded),
+        op.CastLike(scale, query_expanded),
     )
 
     # --- Build Scan for sequential recurrence ---

src/mobius/functions/packed_multi_head_attention.py

@@ -123,8 +123,8 @@ def packed_multi_head_attention() -> ir.Function:
     # Convert to attention bias: 0 for same segment, -10000 for different
     attention_bias = op.Where(
         same_segment,
-        op.Constant(value_float=0.0),
-        op.Constant(value_float=-10000.0),
+        0.0,
+        -10000.0,
     )
     # Reshape for Attention: (1, 1, N, N)
     attention_bias = op.Unsqueeze(attention_bias, [0, 1])

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

@@ -297,15 +297,15 @@ def forward(
         seq_len = op.Shape(input_ids, start=1, end=2)
         _causal_mask = op.Trilu(
             op.Expand(
-                op.Constant(value_float=0.0),
+                0.0,
                 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),
+                -10000.0,
                 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

src/mobius/models/deepseek_ocr2.py

@@ -396,13 +396,13 @@ def forward(
 
         # Cumulative sum to map flat image_features indices
         mask_int = op.Cast(image_mask, to=7)  # INT64
-        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))
 
         # Pad image_features for text-only safety
         pad_row = op.Expand(
-            op.Constant(value_float=0.0),
+            0.0,
             op.Concat(
                 op.Constant(value_ints=[1]),
                 op.Shape(image_features, start=1, end=2),

src/mobius/models/diffllama.py

@@ -175,7 +175,7 @@ def forward(
 
         # Scaled dot-product: Q_4d @ K_4d^T / sqrt(D) → [B, H, S, Sk]
         k_t = op.Transpose(k_4d, perm=[0, 1, 3, 2])  # [B, H, D, Sk]
-        scale = op.Constant(value_float=float(self._scaling))
+        scale = float(self._scaling)
         attn_scores = op.Mul(op.MatMul(q_4d, k_t), scale)
 
         # Causal mask: build [Sk, Sk] lower-triangle via Trilu, then slice
@@ -188,16 +188,16 @@ def forward(
         past_len = op.Sub(kv_len, q_len)  # [1]
         causal_slice = op.Slice(lower_tri, past_len, kv_len, [0])  # [S, Sk]
         causal_bool = op.Cast(causal_slice, to=9)  # BOOL (dtype 9)
-        zero_f = op.Constant(value_float=0.0)
-        neg_inf_f = op.Constant(value_float=float("-inf"))
+        zero_f = 0.0
+        neg_inf_f = float("-inf")
         causal_bias = op.Where(causal_bool, zero_f, neg_inf_f)  # [S, Sk]
         # Unsqueeze to [1, 1, S, Sk] and add to scores [B, H, S, Sk]
         attn_scores = op.Add(attn_scores, op.Unsqueeze(causal_bias, [0, 1]))
 
         # Padding mask (3D bool [B, S, Sk]) → additive bias [B, 1, S, Sk]
         if attention_bias is not None:
-            zero_f2 = op.Constant(value_float=0.0)
-            neg_inf_f2 = op.Constant(value_float=float("-inf"))
+            zero_f2 = 0.0
+            neg_inf_f2 = float("-inf")
             pad_bias = op.Where(attention_bias, zero_f2, neg_inf_f2)
             attn_scores = op.Add(attn_scores, op.Unsqueeze(pad_bias, [1]))
 
@@ -213,7 +213,7 @@ def forward(
         #   lambda_full = exp(lq1·lk1) - exp(lq2·lk2) + lambda_init
         lam1 = op.Exp(op.ReduceSum(op.Mul(self.lambda_q1, self.lambda_k1), keepdims=False))
         lam2 = op.Exp(op.ReduceSum(op.Mul(self.lambda_q2, self.lambda_k2), keepdims=False))
-        lam_init = op.Constant(value_float=float(self._lambda_init))
+        lam_init = float(self._lambda_init)
         lam_full = op.Add(op.Sub(lam1, lam2), lam_init)
 
         # Differential output: out1 - lambda_full * out2  →  [B, H/2, S, 2D]
@@ -222,9 +222,9 @@ def forward(
         # GroupNorm = RMSNorm without learnable affine params (on last dim = 2D)
         #   rms_norm(x) = x / sqrt(mean(x²) + eps)
         x_sq = op.Mul(diff, diff)
-        eps = op.Constant(value_float=float(self._rms_norm_eps))
+        eps = float(self._rms_norm_eps)
         rms = op.Sqrt(op.Add(op.ReduceMean(x_sq, [-1], keepdims=True), eps))
-        scale_norm = op.Constant(value_float=float(1.0 - self._lambda_init))
+        scale_norm = float(1.0 - self._lambda_init)
         normed = op.Mul(op.Div(diff, rms), scale_norm)
 
         # Reshape [B, H/2, S, 2D] → [B, S, H*D]

src/mobius/models/dit.py

@@ -148,7 +148,7 @@ def forward(
         )
 
         # Apply scale and shift to normed hidden_states for self-attention
-        one = op.Constant(value_float=1.0)
+        one = 1.0
         attn_input = op.Mul(normed, op.Add(one, op.Unsqueeze(scale_msa, [1])))
         attn_input = op.Add(attn_input, op.Unsqueeze(shift_msa, [1]))
 

src/mobius/models/falcon.py

@@ -295,8 +295,8 @@ def _create_alibi_bias(op, num_heads: int, seq_len, total_len):
     # Causal mask: mask future positions with large negative value
     causal_mask = op.Where(
         op.GreaterOrEqual(q_with_offset, kv_expanded),
-        op.Constant(value_float=0.0),
-        op.Constant(value_float=-10000.0),
+        0.0,
+        -10000.0,
     )  # [seq_len, total_len]
     causal_4d = op.Unsqueeze(causal_mask, [0, 1])  # [1, 1, seq_len, total_len]
     return op.Add(alibi, causal_4d)

src/mobius/models/flux_sd3.py

@@ -119,7 +119,7 @@ def forward(
             hidden_states,
             temb,
         )
-        one = op.Constant(value_float=1.0)
+        one = 1.0
         img_input = op.Mul(normed, op.Add(one, op.Unsqueeze(scale_msa, [1])))
         img_input = op.Add(img_input, op.Unsqueeze(shift_msa, [1]))
 
@@ -324,7 +324,7 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value, temb: ir.Value
             hidden_states,
             temb,
         )
-        one = op.Constant(value_float=1.0)
+        one = 1.0
         attn_input = op.Mul(normed, op.Add(one, op.Unsqueeze(scale_msa, [1])))
         attn_input = op.Add(attn_input, op.Unsqueeze(shift_msa, [1]))
 

src/mobius/models/gemma3.py

@@ -135,7 +135,7 @@ def forward(self, op: builder.OpBuilder, input_ids: ir.Value, image_features: ir
         image_mask_3d = op.Unsqueeze(image_mask, [-1])
 
         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/granitemoehybrid.py

@@ -120,7 +120,7 @@ def forward(
             op, hidden_states, conv_state, ssm_state
         )
         # residual + output * residual_multiplier
-        rm = op.CastLike(op.Constant(value_float=self._residual_multiplier), mamba_out)
+        rm = op.CastLike(self._residual_multiplier, mamba_out)
         hidden_states = op.Add(residual, op.Mul(mamba_out, rm))
 
         # MoE + shared-MLP path with pre-norm
@@ -131,7 +131,7 @@ def forward(
             self.block_sparse_moe(op, hidden_states),
             self.shared_mlp(op, hidden_states),
         )
-        rm = op.CastLike(op.Constant(value_float=self._residual_multiplier), hidden_states)
+        rm = op.CastLike(self._residual_multiplier, hidden_states)
         hidden_states = op.Add(residual, op.Mul(hidden_states, rm))
 
         return hidden_states, (new_conv_state, new_ssm_state)
@@ -189,7 +189,7 @@ def forward(
             position_embeddings=None,  # NoPE: skip rotary embedding application
             past_key_value=past_key_value,
         )
-        rm = op.CastLike(op.Constant(value_float=self._residual_multiplier), attn_out)
+        rm = op.CastLike(self._residual_multiplier, attn_out)
         hidden_states = op.Add(residual, op.Mul(attn_out, rm))
 
         # MoE + shared-MLP path with pre-norm
@@ -199,7 +199,7 @@ def forward(
             self.block_sparse_moe(op, hidden_states),
             self.shared_mlp(op, hidden_states),
         )
-        rm = op.CastLike(op.Constant(value_float=self._residual_multiplier), hidden_states)
+        rm = op.CastLike(self._residual_multiplier, hidden_states)
         hidden_states = op.Add(residual, op.Mul(hidden_states, rm))
 
         return hidden_states, present_kv

src/mobius/models/hunyuan_dit.py

@@ -109,7 +109,7 @@ def forward(self, op: builder.OpBuilder, x: ir.Value):
         mean = op.ReduceMean(x, [-1], keepdims=True)
         diff = op.Sub(x, mean)
         var = op.ReduceMean(op.Mul(diff, diff), [-1], keepdims=True)
-        eps = op.Constant(value_float=self._eps)
+        eps = self._eps
         return op.Div(diff, op.Sqrt(op.Add(var, eps)))
 
 
@@ -383,7 +383,7 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value, temb: ir.Value
         scale, shift = op.Split(emb, num_outputs=2, axis=-1, _outputs=2)
         # Norm → (1 + scale) * normed + shift
         normed = self.norm(op, hidden_states)
-        one = op.Constant(value_float=1.0)
+        one = 1.0
         normed = op.Mul(normed, op.Add(one, op.Unsqueeze(scale, [1])))
         return op.Add(normed, op.Unsqueeze(shift, [1]))