Fix integer true-divide silently truncating instead of promoting to float

coreai_torch/_aten_to_core.py

@@ -810,8 +810,6 @@ def replace_binary_ops(
         "add.Tensor": coreai.broadcasting_add,
         "add.Scalar": coreai.broadcasting_add,
         "add": coreai.broadcasting_add,
-        "div.Tensor": coreai.broadcasting_divide,
-        "div.Scalar": coreai.broadcasting_divide,
         "maximum.default": coreai.broadcasting_maximum,
         "minimum.default": coreai.broadcasting_minimum,
         "fmod.Tensor": coreai.broadcasting_modulo,
@@ -876,13 +874,20 @@ def replace_div_tensor_mode(
         else (node.args[2] if len(node.args) > 2 else None)
     )
 
+    if rounding_mode is None:
+        # True divide: integer operands must promote to the node's float
+        # output type before dividing, not the generic same-kind-stays-integer
+        # promotion rule used by "floor"/"trunc" below.
+        result_type = get_output_element_type_from_node(node)
+        return coreai.broadcasting_divide(
+            coreai.cast(x, result_type), coreai.cast(y, result_type)
+        )
+
     promoted_type = get_promoted_type(x.type, y.type)
     casted_x = coreai.cast(x, promoted_type)
     casted_y = coreai.cast(y, promoted_type)
 
-    if rounding_mode is None:
-        return coreai.broadcasting_divide(casted_x, casted_y)
-    elif rounding_mode == "floor":
+    if rounding_mode == "floor":
         return coreai.broadcasting_floor_divide(casted_x, casted_y)
     elif rounding_mode == "trunc":
         # Integer division already truncates toward zero, so a plain divide
@@ -3588,8 +3593,8 @@ def sdpa_maskless(q: Value, k: Value, v: Value) -> Value:
     "cos.default": replace_unary_ops,
     "cosh.default": replace_unary_ops,
     "cumsum.default": replace_cumsum,
-    "div.Scalar": replace_binary_ops,
-    "div.Tensor": replace_binary_ops,
+    "div.Scalar": replace_truediv,
+    "div.Tensor": replace_truediv,
     "div.Tensor_mode": replace_div_tensor_mode,
     "embedding.default": replace_embedding,
     "empty.default": replace_empty,
@@ -3719,7 +3724,7 @@ def sdpa_maskless(q: Value, k: Value, v: Value) -> Value:
     "truediv": replace_truediv,
     "to.dtype": replace_to_dtype,
     "topk.default": replace_topk,
-    "true_divide.Tensor": replace_binary_ops,
+    "true_divide.Tensor": replace_truediv,
     "trunc.default": replace_trunc,
     "trunc": replace_trunc,
     "unsqueeze.default": replace_unsqueeze,

docs/api/TorchConverter.md

@@ -241,16 +241,20 @@ import torch
 from coreai._compiler.dialects import coreai
 from coreai_torch._utils import get_operands
 
+
 @torch.library.custom_op("my_lib::scaled_add", mutates_args=())
 def scaled_add(x: torch.Tensor, y: torch.Tensor, scale: float) -> torch.Tensor:
     return x + scale * y
 
+
 @scaled_add.register_fake
 def _(x, y, scale):
     return torch.empty_like(x)
 
+
 converter = TorchConverter()
 
+
 @converter.register_torch_lowering("my_lib::scaled_add.default")
 def lower_scaled_add(values_map, node, loc):
     x, y = get_operands(values_map, node, [0, 1], loc)
@@ -259,6 +263,7 @@ def lower_scaled_add(values_map, node, loc):
     scaled_y = coreai.broadcasting_mul(y, scale_val, loc=loc)
     return coreai.broadcasting_add(x, scaled_y, loc=loc)
 
+
 coreai_program = converter.add_exported_program(exported).to_coreai()
 coreai_program.optimize()
 ```
@@ -272,7 +277,10 @@ from coreai_torch._utils import get_operand
 
 converter = TorchConverter()
 
-@converter.register_torch_lowering("aten::_adaptive_avg_pool2d.default", allow_override=True)
+
+@converter.register_torch_lowering(
+    "aten::_adaptive_avg_pool2d.default", allow_override=True
+)
 def lower_adaptive_avg_pool2d_static(values_map, node, loc):
     x = get_operand(values_map, node, 0, loc)
     output_h, output_w = node.args[1]
@@ -290,6 +298,7 @@ def lower_adaptive_avg_pool2d_static(values_map, node, loc):
         coreai.cast(float(kernel_h * kernel_w), x.type.element_type),
     )
 
+
 coreai_program = converter.add_exported_program(exported).to_coreai()
 coreai_program.optimize()
 ```
@@ -317,7 +326,12 @@ Registers one or more `TorchMetalKernel` objects so the converter can convert th
 
 ```python
 import torch
-from coreai_torch import TorchConverter, TorchMetalKernel, MetalParameter, get_decomp_table
+from coreai_torch import (
+    TorchConverter,
+    TorchMetalKernel,
+    MetalParameter,
+    get_decomp_table,
+)
 
 
 def torch_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
@@ -383,9 +397,9 @@ coreai_program = (
     TorchConverter()
     .add_pytorch_module(
         model,
-        export_fn=lambda m: torch.export.export(m, args=example_inputs).run_decompositions(
-            coreai_torch.get_decomp_table()
-        ),
+        export_fn=lambda m: torch.export.export(
+            m, args=example_inputs
+        ).run_decompositions(coreai_torch.get_decomp_table()),
     )
     .to_coreai()
 )
@@ -457,6 +471,7 @@ class Linear(nn.Module):
     def forward(self, x):
         return self.fc(x)
 
+
 ep = torch.export.export(Linear().eval(), args=(torch.randn(1, 8),))
 ep = ep.run_decompositions(get_decomp_table())
 
@@ -473,16 +488,17 @@ TorchConverter().add_exported_program(
 class KVCache(nn.Module):
     def __init__(self):
         super().__init__()
-        self.register_buffer("kv_cache", torch.zeros(1, 4))   # state[0]
-        self.register_buffer("pos_idx", torch.zeros(1))       # state[1]
+        self.register_buffer("kv_cache", torch.zeros(1, 4))  # state[0]
+        self.register_buffer("pos_idx", torch.zeros(1))  # state[1]
 
     def forward(self, x, y, z):
-        self.kv_cache.add_(x)       # buffer mutation
-        self.pos_idx.add_(1)        # buffer mutation
-        y.mul_(2)                   # state[2]: mutated user input
+        self.kv_cache.add_(x)  # buffer mutation
+        self.pos_idx.add_(1)  # buffer mutation
+        y.mul_(2)  # state[2]: mutated user input
         # non-mutated: x -> input[0], z -> input[1]
         return self.kv_cache + y, z * 3
 
+
 ep = torch.export.export(
     KVCache().eval(),
     args=(torch.randn(1, 4), torch.randn(1, 4), torch.randn(1, 4)),

docs/api/composite-ops.md

@@ -7,7 +7,13 @@ Use these module subclasses (and ATen-derived ops) to preserve an operation's bo
 **Public import:**
 
 ```python
-from coreai_torch.composite_ops import GatherMM, GatedDeltaUpdate, RMSNormImpl, RoPE, SDPA
+from coreai_torch.composite_ops import (
+    GatherMM,
+    GatedDeltaUpdate,
+    RMSNormImpl,
+    RoPE,
+    SDPA,
+)
 ```
 
 coreai-torch provides composite ops in two categories:

docs/api/composite-ops/gather-mm.md

@@ -82,10 +82,10 @@ class MoELayer(nn.Module):
 
     def forward(
         self,
-        x: torch.Tensor,          # [B, T, 1, 1, D]
-        experts: torch.Tensor,    # [E, D, H]
-        indices: torch.Tensor,    # [B, T, K]
-    ) -> torch.Tensor:            # [B, T, K, 1, H]
+        x: torch.Tensor,  # [B, T, 1, 1, D]
+        experts: torch.Tensor,  # [E, D, H]
+        indices: torch.Tensor,  # [B, T, K]
+    ) -> torch.Tensor:  # [B, T, K, 1, H]
         return self.gather_mm(x, experts, rhs_indices=indices)
 
 
@@ -129,10 +129,10 @@ class FusedMoELayer(nn.Module):
 
     def forward(
         self,
-        x: torch.Tensor,            # [B, T, 1, 1, D]
-        fused_experts: torch.Tensor, # [2, E, D, H]  (gate + up stacked)
-        indices: torch.Tensor,      # [B, T, K]
-    ) -> torch.Tensor:              # [2, B, T, K, 1, H]
+        x: torch.Tensor,  # [B, T, 1, 1, D]
+        fused_experts: torch.Tensor,  # [2, E, D, H]  (gate + up stacked)
+        indices: torch.Tensor,  # [B, T, K]
+    ) -> torch.Tensor:  # [2, B, T, K, 1, H]
         return self.gather_mm(x, fused_experts, rhs_indices=indices)
 ```
 
@@ -169,10 +169,11 @@ def _gather(x, indices, num_batch_axes=0):
     flat_indices = indices.to(torch.int32).flatten()
     flat_gather = torch.index_select(x, dim=num_batch_axes, index=flat_indices)
     result_shape = (
-        x.shape[:num_batch_axes] + indices.shape + x.shape[num_batch_axes + 1:]
+        x.shape[:num_batch_axes] + indices.shape + x.shape[num_batch_axes + 1 :]
     )
     return flat_gather.view(result_shape)
 
+
 def gather_mm(lhs, rhs, lhs_indices=None, rhs_indices=None, num_batch_axes=0):
     if lhs_indices is not None:
         lhs = _gather(lhs, lhs_indices, num_batch_axes=num_batch_axes)

docs/api/composite-ops/instance-norm.md

@@ -40,12 +40,15 @@ gamma = torch.randn(C)
 beta = torch.randn(C)
 
 output = torch.ops.aten.instance_norm.default(
-    input, gamma, beta,
-    None, None,        # running_mean / running_var unused in inference
-    True,              # use_input_stats
-    0.1,               # momentum (ignored in inference)
-    1e-5,              # eps
-    True,              # cudnn_enabled (ignored)
+    input,
+    gamma,
+    beta,
+    None,
+    None,  # running_mean / running_var unused in inference
+    True,  # use_input_stats
+    0.1,  # momentum (ignored in inference)
+    1e-5,  # eps
+    True,  # cudnn_enabled (ignored)
 )
 ```
 

docs/api/debugging.md

@@ -84,14 +84,12 @@ from coreai_torch.debugging.comparator import create_comparator_for_programs
 comparator = await create_comparator_for_programs(
     source_program=exported_program,
     target_program=coreai_program,
-    target_entry_point="main"
+    target_entry_point="main",
 )
 
 # Compare outputs with tolerance
 result = await comparator.compare_with_tolerance(
-    inputs={"x": example_input},
-    rtol=1e-5,
-    atol=1e-8
+    inputs={"x": example_input}, rtol=1e-5, atol=1e-8
 )
 
 # Check for differences
@@ -123,8 +121,7 @@ coreai_op_ids = [1, 5, 10, 15]
 
 # Capture intermediate values
 results = await inspector.get_intermediates_for_ops(
-    coreai_op_ids,
-    inputs={"x": np.random.randn(2, 4).astype(np.float32)}
+    coreai_op_ids, inputs={"x": np.random.randn(2, 4).astype(np.float32)}
 )
 
 # Check results
@@ -144,7 +141,7 @@ Analyze structural differences between model implementations using graph isomorp
 from coreai_torch.debugging.graph_diff import (
     compute_exported_program_diff,
     compute_coreai_program_diff,
-    write_diff
+    write_diff,
 )
 
 # Compare two PyTorch programs
@@ -160,12 +157,7 @@ else:
     print(f"✗ Found {diff.summary.unmapped_source_node_count} structural differences")
 
     # Write detailed diff report to stdout
-    write_diff(
-        diff,
-        diff.source_graph,
-        diff.target_graph,
-        max_items=20
-    )
+    write_diff(diff, diff.source_graph, diff.target_graph, max_items=20)
 ```
 
 
@@ -180,9 +172,7 @@ from coreai_torch.debugging.benchmarker import benchmark_coreai_program
 
 # Run benchmark
 result = await benchmark_coreai_program(
-    coreai_program=coreai_program,
-    inputs={"x": torch.randn(2, 4)},
-    num_runs=50
+    coreai_program=coreai_program, inputs={"x": torch.randn(2, 4)}, num_runs=50
 )
 
 # Show timing summary
@@ -203,10 +193,8 @@ Create custom checks beyond NaN/infinity:
 ```python
 def check_large_values(outputs):
     """Check if any output has values > threshold"""
-    return any(
-        abs(arr).max() > 1000.0 if arr is not None else False
-        for arr in outputs
-    )
+    return any(abs(arr).max() > 1000.0 if arr is not None else False for arr in outputs)
+
 
 # Use custom check
 result = await validator.check(check_large_values, inputs=example_input)
@@ -256,9 +244,7 @@ exported_program = torch.export.export(model, args=example_input)
 
 # Save intermediate values to disk
 metadata_path = save_intermediates(
-    program=exported_program,
-    inputs=example_input,
-    output_dir=Path("./debug_output")
+    program=exported_program, inputs=example_input, output_dir=Path("./debug_output")
 )
 
 print(f"Intermediates saved to: {metadata_path}")
@@ -291,12 +277,13 @@ def custom_filter(node, result):
     """Only save convolution and linear layer outputs"""
     return any(op in str(node.target).lower() for op in ["conv", "linear", "matmul"])
 
+
 # Save only filtered operations
 metadata_path = save_intermediates(
     program=exported_program,
     inputs=example_input,
     output_dir=Path("./debug_output"),
-    node_filter=custom_filter
+    node_filter=custom_filter,
 )
 ```
 

docs/getting-started/installation.md

@@ -46,6 +46,7 @@ Run the following to confirm coreai-torch is installed correctly — a version s
 
 ```python
 import coreai_torch
+
 print(coreai_torch.__version__)
 ```