From 629cce3baa188dc72dfb49c886c70f3ffb87f5ef Mon Sep 17 00:00:00 2001 From: gokulkrishna98 Date: Tue, 7 Jul 2026 14:37:13 -0700 Subject: [PATCH 1/4] Fix integer true-divide silently truncating instead of promoting to float aten.div.Tensor, div.Scalar, and true_divide.Tensor on integer operands were dispatched to the generic replace_binary_ops handler, which divides using the generic same-kind-stays-integer promotion rule and only casts the already-truncated result to float afterward. True divide promotes integer operands to a float dtype before dividing, so this silently dropped the fractional part on every backend. Re-point these ops at replace_truediv, which already promotes operands to the correct float result type before dividing. --- coreai_torch/_aten_to_core.py | 8 +++----- tests/ops/test_ops.py | 35 +++++++++++++++++++++++++++++++++++ 2 files changed, 38 insertions(+), 5 deletions(-) diff --git a/coreai_torch/_aten_to_core.py b/coreai_torch/_aten_to_core.py index 1f292e7..d825df8 100644 --- a/coreai_torch/_aten_to_core.py +++ b/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, @@ -3488,8 +3486,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, @@ -3619,7 +3617,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, diff --git a/tests/ops/test_ops.py b/tests/ops/test_ops.py index a1daf19..4d36c57 100644 --- a/tests/ops/test_ops.py +++ b/tests/ops/test_ops.py @@ -1541,6 +1541,41 @@ def forward(self, x: Tensor, y: Tensor) -> Tensor: ) +@pytest.mark.parametrize( + "x,y", + [ + ( + torch.tensor([7, -7, 3, 1], dtype=torch.int32), + torch.tensor([2, 2, 2, 4], dtype=torch.int32), + ), + ( + torch.tensor([1, 2, 3, 4], dtype=torch.int64), + torch.tensor([3, 3, 3, 3], dtype=torch.int64), + ), + ], +) +async def test_div_integer_promotes_to_float(x: Tensor, y: Tensor) -> None: + """aten.div.Tensor on integer operands must promote to float before dividing.""" + + class DivModel(nn.Module): + def forward(self, x: Tensor, y: Tensor) -> Tensor: + return x / y + + model = DivModel().eval() + await validate_numerical_output(model=model, x=x, y=y) + + +async def test_div_scalar_integer_promotes_to_float() -> None: + """aten.div.Scalar on an integer tensor must promote to float before dividing.""" + x = torch.tensor([7, -7, 3, 1], dtype=torch.int32) + + class DivScalarModel(nn.Module): + def forward(self, x: Tensor) -> Tensor: + return x / 4 + + await validate_numerical_output(model=DivScalarModel().eval(), x=x) + + @pytest.mark.parametrize( "x,y", [ From 25082e40c9dc055ab8b3939f4669e0200fbef8c3 Mon Sep 17 00:00:00 2001 From: gokulkrishna98 Date: Tue, 7 Jul 2026 15:13:27 -0700 Subject: [PATCH 2/4] address PR review feedback Fix adjacent latent same-class bug: replace_div_tensor_mode with rounding_mode=None used the generic integer-stays-integer promotion rule instead of promoting to the node output float type before dividing. Add torch.true_divide(int, int) test to directly cover the true_divide.Tensor resolver rewiring. Add torch.div(int, int, rounding_mode=None) test to cover the div_tensor_mode fix. --- coreai_torch/_aten_to_core.py | 13 ++++++++++--- tests/ops/test_ops.py | 25 +++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 3 deletions(-) diff --git a/coreai_torch/_aten_to_core.py b/coreai_torch/_aten_to_core.py index d825df8..65a22fa 100644 --- a/coreai_torch/_aten_to_core.py +++ b/coreai_torch/_aten_to_core.py @@ -874,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 diff --git a/tests/ops/test_ops.py b/tests/ops/test_ops.py index 4d36c57..b4527ef 100644 --- a/tests/ops/test_ops.py +++ b/tests/ops/test_ops.py @@ -1576,6 +1576,31 @@ def forward(self, x: Tensor) -> Tensor: await validate_numerical_output(model=DivScalarModel().eval(), x=x) +async def test_true_divide_integer_promotes_to_float() -> None: + """aten.true_divide.Tensor on integer operands must promote to float before dividing.""" + x = torch.tensor([7, -7, 3, 1], dtype=torch.int32) + y = torch.tensor([2, 2, 2, 4], dtype=torch.int32) + + class TrueDivideModel(nn.Module): + def forward(self, x: Tensor, y: Tensor) -> Tensor: + return torch.true_divide(x, y) + + await validate_numerical_output(model=TrueDivideModel().eval(), x=x, y=y) + + +async def test_div_tensor_mode_none_integer_promotes_to_float() -> None: + """aten.div.Tensor_mode with rounding_mode=None on integer operands must + promote to float before dividing, matching aten.div.Tensor semantics.""" + x = torch.tensor([7, -7, 3, 1], dtype=torch.int32) + y = torch.tensor([2, 2, 2, 4], dtype=torch.int32) + + class DivTensorModeModel(nn.Module): + def forward(self, x: Tensor, y: Tensor) -> Tensor: + return torch.div(x, y, rounding_mode=None) + + await validate_numerical_output(model=DivTensorModeModel().eval(), x=x, y=y) + + @pytest.mark.parametrize( "x,y", [ From fbc1b474838ded85d5f060b2335c60d6c7045834 Mon Sep 17 00:00:00 2001 From: gokulkrishna98 Date: Tue, 7 Jul 2026 15:28:20 -0700 Subject: [PATCH 3/4] Group div-family tests into a TestDiv class Combines test_div, test_div_integer_promotes_to_float, test_div_scalar_integer_promotes_to_float, test_true_divide_integer_promotes_to_float, test_div_tensor_mode_none_integer_promotes_to_float, test_div_tensor_mode, test_true_divide, and test_true_divide_scalar into a single TestDiv class, matching the per-op class convention used elsewhere in the file (e.g. TestCopy). --- tests/ops/test_ops.py | 269 +++++++++++++++++++++--------------------- 1 file changed, 134 insertions(+), 135 deletions(-) diff --git a/tests/ops/test_ops.py b/tests/ops/test_ops.py index b4527ef..8b8a52d 100644 --- a/tests/ops/test_ops.py +++ b/tests/ops/test_ops.py @@ -1519,166 +1519,165 @@ def forward(self, dest: Tensor, src: Tensor) -> Tensor: ) -@pytest.mark.parametrize("dynamic", [False, True]) -@pytest.mark.parametrize("x", [torch.rand(2, 2)]) -@pytest.mark.parametrize("y", [torch.rand(2, 2)]) -async def test_div(x: Tensor, y: Tensor, dynamic: bool) -> None: - class DivModel(nn.Module): - def __init__(self) -> None: - super().__init__() - - def forward(self, x: Tensor, y: Tensor) -> Tensor: - return x / y - - model = DivModel().eval() - if dynamic: - dims = _all_dims_dynamic(x) - dynamic_shapes = {"x": dims, "y": dims} - else: - dynamic_shapes = None - await validate_numerical_output( - model=model, x=x, y=y, dynamic_shapes=dynamic_shapes - ) - +class TestDiv: + """Test suite for aten.div.Tensor / div.Scalar / div.Tensor_mode / + true_divide.Tensor → coreai.broadcasting_divide conversion.""" -@pytest.mark.parametrize( - "x,y", - [ - ( - torch.tensor([7, -7, 3, 1], dtype=torch.int32), - torch.tensor([2, 2, 2, 4], dtype=torch.int32), - ), - ( - torch.tensor([1, 2, 3, 4], dtype=torch.int64), - torch.tensor([3, 3, 3, 3], dtype=torch.int64), - ), - ], -) -async def test_div_integer_promotes_to_float(x: Tensor, y: Tensor) -> None: - """aten.div.Tensor on integer operands must promote to float before dividing.""" - - class DivModel(nn.Module): - def forward(self, x: Tensor, y: Tensor) -> Tensor: - return x / y + @pytest.mark.parametrize("dynamic", [False, True]) + @pytest.mark.parametrize("x", [torch.rand(2, 2)]) + @pytest.mark.parametrize("y", [torch.rand(2, 2)]) + async def test_div(self, x: Tensor, y: Tensor, dynamic: bool) -> None: + class DivModel(nn.Module): + def __init__(self) -> None: + super().__init__() - model = DivModel().eval() - await validate_numerical_output(model=model, x=x, y=y) + def forward(self, x: Tensor, y: Tensor) -> Tensor: + return x / y + model = DivModel().eval() + if dynamic: + dims = _all_dims_dynamic(x) + dynamic_shapes = {"x": dims, "y": dims} + else: + dynamic_shapes = None + await validate_numerical_output( + model=model, x=x, y=y, dynamic_shapes=dynamic_shapes + ) -async def test_div_scalar_integer_promotes_to_float() -> None: - """aten.div.Scalar on an integer tensor must promote to float before dividing.""" - x = torch.tensor([7, -7, 3, 1], dtype=torch.int32) + @pytest.mark.parametrize( + "x,y", + [ + ( + torch.tensor([7, -7, 3, 1], dtype=torch.int32), + torch.tensor([2, 2, 2, 4], dtype=torch.int32), + ), + ( + torch.tensor([1, 2, 3, 4], dtype=torch.int64), + torch.tensor([3, 3, 3, 3], dtype=torch.int64), + ), + ], + ) + async def test_div_integer_promotes_to_float(self, x: Tensor, y: Tensor) -> None: + """aten.div.Tensor on integer operands must promote to float before dividing.""" - class DivScalarModel(nn.Module): - def forward(self, x: Tensor) -> Tensor: - return x / 4 + class DivModel(nn.Module): + def forward(self, x: Tensor, y: Tensor) -> Tensor: + return x / y - await validate_numerical_output(model=DivScalarModel().eval(), x=x) + model = DivModel().eval() + await validate_numerical_output(model=model, x=x, y=y) + async def test_div_scalar_integer_promotes_to_float(self) -> None: + """aten.div.Scalar on an integer tensor must promote to float before dividing.""" + x = torch.tensor([7, -7, 3, 1], dtype=torch.int32) -async def test_true_divide_integer_promotes_to_float() -> None: - """aten.true_divide.Tensor on integer operands must promote to float before dividing.""" - x = torch.tensor([7, -7, 3, 1], dtype=torch.int32) - y = torch.tensor([2, 2, 2, 4], dtype=torch.int32) + class DivScalarModel(nn.Module): + def forward(self, x: Tensor) -> Tensor: + return x / 4 - class TrueDivideModel(nn.Module): - def forward(self, x: Tensor, y: Tensor) -> Tensor: - return torch.true_divide(x, y) + await validate_numerical_output(model=DivScalarModel().eval(), x=x) - await validate_numerical_output(model=TrueDivideModel().eval(), x=x, y=y) + async def test_true_divide_integer_promotes_to_float(self) -> None: + """aten.true_divide.Tensor on integer operands must promote to float before dividing.""" + x = torch.tensor([7, -7, 3, 1], dtype=torch.int32) + y = torch.tensor([2, 2, 2, 4], dtype=torch.int32) + class TrueDivideModel(nn.Module): + def forward(self, x: Tensor, y: Tensor) -> Tensor: + return torch.true_divide(x, y) -async def test_div_tensor_mode_none_integer_promotes_to_float() -> None: - """aten.div.Tensor_mode with rounding_mode=None on integer operands must - promote to float before dividing, matching aten.div.Tensor semantics.""" - x = torch.tensor([7, -7, 3, 1], dtype=torch.int32) - y = torch.tensor([2, 2, 2, 4], dtype=torch.int32) + await validate_numerical_output(model=TrueDivideModel().eval(), x=x, y=y) - class DivTensorModeModel(nn.Module): - def forward(self, x: Tensor, y: Tensor) -> Tensor: - return torch.div(x, y, rounding_mode=None) + async def test_div_tensor_mode_none_integer_promotes_to_float(self) -> None: + """aten.div.Tensor_mode with rounding_mode=None on integer operands must + promote to float before dividing, matching aten.div.Tensor semantics.""" + x = torch.tensor([7, -7, 3, 1], dtype=torch.int32) + y = torch.tensor([2, 2, 2, 4], dtype=torch.int32) - await validate_numerical_output(model=DivTensorModeModel().eval(), x=x, y=y) + class DivTensorModeModel(nn.Module): + def forward(self, x: Tensor, y: Tensor) -> Tensor: + return torch.div(x, y, rounding_mode=None) + await validate_numerical_output(model=DivTensorModeModel().eval(), x=x, y=y) -@pytest.mark.parametrize( - "x,y", - [ - # Float tensors - mixed positive/negative values - ( - torch.tensor([[3.5, -7.2], [-2.8, 9.1]]), - torch.tensor([[2.0, 3.0], [2.0, -4.0]]), - ), - # Larger tensors - ( - torch.rand(3, 4) * 10 - 5, - torch.rand(3, 4) * 4 + 0.5, - ), # Avoid division by values near zero - # Broadcasting case - (torch.rand(2, 3, 4) * 10 - 5, torch.rand(1, 3, 1) * 4 + 0.5), - ], -) -@pytest.mark.parametrize("rounding_mode", [None, "floor", "trunc"]) -async def test_div_tensor_mode(x: Tensor, y: Tensor, rounding_mode: str | None) -> None: - """Test division with different rounding modes. - - aten.div.Tensor_mode(input, other, rounding_mode) supports: - - None: True division (standard floating-point division) - - "floor": Floor division (rounds toward negative infinity) - - "trunc": Truncated division (rounds toward zero) - """ - - class DivTensorModeModel(nn.Module): - def __init__(self) -> None: - super().__init__() + @pytest.mark.parametrize( + "x,y", + [ + # Float tensors - mixed positive/negative values + ( + torch.tensor([[3.5, -7.2], [-2.8, 9.1]]), + torch.tensor([[2.0, 3.0], [2.0, -4.0]]), + ), + # Larger tensors + ( + torch.rand(3, 4) * 10 - 5, + torch.rand(3, 4) * 4 + 0.5, + ), # Avoid division by values near zero + # Broadcasting case + (torch.rand(2, 3, 4) * 10 - 5, torch.rand(1, 3, 1) * 4 + 0.5), + ], + ) + @pytest.mark.parametrize("rounding_mode", [None, "floor", "trunc"]) + async def test_div_tensor_mode( + self, x: Tensor, y: Tensor, rounding_mode: str | None + ) -> None: + """Test division with different rounding modes. - def forward(self, x: Tensor, y: Tensor) -> Tensor: - return torch.div(x, y, rounding_mode=rounding_mode) + aten.div.Tensor_mode(input, other, rounding_mode) supports: + - None: True division (standard floating-point division) + - "floor": Floor division (rounds toward negative infinity) + - "trunc": Truncated division (rounds toward zero) + """ - model = DivTensorModeModel().eval() - await validate_numerical_output(model=model, x=x, y=y) + class DivTensorModeModel(nn.Module): + def __init__(self) -> None: + super().__init__() + def forward(self, x: Tensor, y: Tensor) -> Tensor: + return torch.div(x, y, rounding_mode=rounding_mode) -@pytest.mark.parametrize("dynamic", [False, True]) -@pytest.mark.parametrize( - "x,y", - [ - (torch.rand(2, 3) + 0.1, torch.rand(2, 3) + 0.1), - (torch.rand(3, 4, 5) + 0.1, torch.rand(3, 4, 5) + 0.1), - (torch.rand(4) + 0.1, torch.rand(4) + 0.1), - # FP16 - ( - torch.rand(2, 3, dtype=torch.float16) + 0.1, - torch.rand(2, 3, dtype=torch.float16) + 0.1, - ), - ], -) -async def test_true_divide(x: Tensor, y: Tensor, dynamic: bool) -> None: - class TrueDivideModel(nn.Module): - def forward(self, x: Tensor, y: Tensor) -> Tensor: - return torch.true_divide(x, y) + model = DivTensorModeModel().eval() + await validate_numerical_output(model=model, x=x, y=y) - model = TrueDivideModel().eval() - if dynamic: - dims = _all_dims_dynamic(x) - dynamic_shapes = {"x": dims, "y": dims} - else: - dynamic_shapes = None - await validate_numerical_output( - model=model, x=x, y=y, dynamic_shapes=dynamic_shapes + @pytest.mark.parametrize("dynamic", [False, True]) + @pytest.mark.parametrize( + "x,y", + [ + (torch.rand(2, 3) + 0.1, torch.rand(2, 3) + 0.1), + (torch.rand(3, 4, 5) + 0.1, torch.rand(3, 4, 5) + 0.1), + (torch.rand(4) + 0.1, torch.rand(4) + 0.1), + # FP16 + ( + torch.rand(2, 3, dtype=torch.float16) + 0.1, + torch.rand(2, 3, dtype=torch.float16) + 0.1, + ), + ], ) + async def test_true_divide(self, x: Tensor, y: Tensor, dynamic: bool) -> None: + class TrueDivideModel(nn.Module): + def forward(self, x: Tensor, y: Tensor) -> Tensor: + return torch.true_divide(x, y) + model = TrueDivideModel().eval() + if dynamic: + dims = _all_dims_dynamic(x) + dynamic_shapes = {"x": dims, "y": dims} + else: + dynamic_shapes = None + await validate_numerical_output( + model=model, x=x, y=y, dynamic_shapes=dynamic_shapes + ) -@pytest.mark.parametrize("dynamic", [False, True]) -@pytest.mark.parametrize("x", [torch.rand(2, 3) + 0.1, torch.rand(3, 4, 5) + 0.1]) -async def test_true_divide_scalar(x: Tensor, dynamic: bool) -> None: - class TrueDivideScalarModel(nn.Module): - def forward(self, x: Tensor) -> Tensor: - return torch.true_divide(x, 2.0) + @pytest.mark.parametrize("dynamic", [False, True]) + @pytest.mark.parametrize("x", [torch.rand(2, 3) + 0.1, torch.rand(3, 4, 5) + 0.1]) + async def test_true_divide_scalar(self, x: Tensor, dynamic: bool) -> None: + class TrueDivideScalarModel(nn.Module): + def forward(self, x: Tensor) -> Tensor: + return torch.true_divide(x, 2.0) - model = TrueDivideScalarModel().eval() - dynamic_shapes = {"x": _all_dims_dynamic(x)} if dynamic else None - await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes) + model = TrueDivideScalarModel().eval() + dynamic_shapes = {"x": _all_dims_dynamic(x)} if dynamic else None + await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes) @pytest.mark.parametrize("dynamic", [False, True]) From be7f176b5d7385bc8cfe876b9cf7fbc2098ad142 Mon Sep 17 00:00:00 2001 From: gokulkrishna98 Date: Wed, 29 Jul 2026 12:08:46 -0700 Subject: [PATCH 4/4] style: format Python code blocks in docs Markdown ruff 0.16 promoted Markdown code-block formatting out of preview, so CI (which resolves the newest "ruff>=0.12.0") now checks fenced Python blocks in docs/. Reformat the six affected files; no prose or behavior changes. --- docs/api/TorchConverter.md | 36 ++++++++++++++++++------- docs/api/composite-ops.md | 8 +++++- docs/api/composite-ops/gather-mm.md | 19 ++++++------- docs/api/composite-ops/instance-norm.md | 15 ++++++----- docs/api/debugging.md | 35 ++++++++---------------- docs/getting-started/installation.md | 1 + 6 files changed, 64 insertions(+), 50 deletions(-) diff --git a/docs/api/TorchConverter.md b/docs/api/TorchConverter.md index 9ebc704..17c187e 100644 --- a/docs/api/TorchConverter.md +++ b/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)), diff --git a/docs/api/composite-ops.md b/docs/api/composite-ops.md index 3b277e6..23e632c 100644 --- a/docs/api/composite-ops.md +++ b/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: diff --git a/docs/api/composite-ops/gather-mm.md b/docs/api/composite-ops/gather-mm.md index 0dd2f44..5100495 100644 --- a/docs/api/composite-ops/gather-mm.md +++ b/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) diff --git a/docs/api/composite-ops/instance-norm.md b/docs/api/composite-ops/instance-norm.md index 3ddb3a1..f19dc4d 100644 --- a/docs/api/composite-ops/instance-norm.md +++ b/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) ) ``` diff --git a/docs/api/debugging.md b/docs/api/debugging.md index 18a623b..1ffc997 100644 --- a/docs/api/debugging.md +++ b/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, ) ``` diff --git a/docs/getting-started/installation.md b/docs/getting-started/installation.md index 362b5af..d50f26b 100644 --- a/docs/getting-started/installation.md +++ b/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__) ```