From 956e943a51f9279e602bd9eb6b18b74f8aa5f5f9 Mon Sep 17 00:00:00 2001 From: Bharath Vedartham Date: Tue, 25 Aug 2026 14:20:31 +0530 Subject: [PATCH 1/2] add trunc op for mlx backend --- backends/mlx/ops.py | 5 ++--- backends/mlx/runtime/MLXInterpreter.h | 8 ++++++++ backends/mlx/serialization/schema.fbs | 8 +++++++- backends/mlx/test/test_ops.py | 1 + 4 files changed, 18 insertions(+), 4 deletions(-) diff --git a/backends/mlx/ops.py b/backends/mlx/ops.py index 62378f54c07..010b1b35b85 100644 --- a/backends/mlx/ops.py +++ b/backends/mlx/ops.py @@ -159,6 +159,7 @@ TransposeNode, TrilNode, TriuNode, + TruncNode, UpdateAndAttendNode, VarNode, VidOrTid, @@ -384,6 +385,7 @@ def normalize_reduction_dim( # Rounding (torch.ops.aten.floor.default, FloorNode, "aten.floor"), (torch.ops.aten.ceil.default, CeilNode, "aten.ceil"), + (torch.ops.aten.trunc.default, TruncNode, "aten.trunc"), # Powers / roots (torch.ops.aten.square.default, SquareNode, "aten.square"), (torch.ops.aten.exp.default, ExpNode, "aten.exp"), @@ -836,7 +838,6 @@ def handler(P: MLXProgramBuilder, n: Node) -> Slot: _make_binary_handler(_node_cls, _op_name, _lift_b) ) - _SCALAR_INT_OPS: List[Tuple[Any, Any, str]] = [ (operator.add, AddIntNode, "operator.add"), (operator.sub, SubtractIntNode, "operator.sub"), @@ -937,7 +938,6 @@ def handler(P: MLXProgramBuilder, n: Node) -> Slot: _make_reduction_handler(_node_cls, _op_name, _max_args) ) - _FULL_OPS: List[Tuple[List[Any], str, Optional[float]]] = [ ([torch.ops.aten.full.default], "aten.full", None), ([torch.ops.aten.zeros.default], "aten.zeros", 0.0), @@ -988,7 +988,6 @@ def handler(P: MLXProgramBuilder, n: Node) -> Slot: for _targets, _op_name, _fixed_fill in _FULL_OPS: REGISTRY.register(target=_targets)(_make_full_handler(_op_name, _fixed_fill)) - _FULL_LIKE_OPS: List[Tuple[List[Any], str, Optional[float]]] = [ ([torch.ops.aten.full_like.default], "aten.full_like", None), ([torch.ops.aten.zeros_like.default], "aten.zeros_like", 0.0), diff --git a/backends/mlx/runtime/MLXInterpreter.h b/backends/mlx/runtime/MLXInterpreter.h index b14fdf6903c..4f0087de464 100644 --- a/backends/mlx/runtime/MLXInterpreter.h +++ b/backends/mlx/runtime/MLXInterpreter.h @@ -1545,6 +1545,11 @@ inline void exec_ceil(const CeilNode& n, ExecutionState& st, StreamOrDevice s) { st.set_tensor(n.out, ceil(st.const_tensor_ref(n.x), s)); } +inline void +exec_trunc(const TruncNode& n, ExecutionState& st, StreamOrDevice s) { + st.set_tensor(n.out, trunc(st.const_tensor_ref(n.x), s)); +} + inline void exec_square(const SquareNode& n, ExecutionState& st, StreamOrDevice s) { st.set_tensor(n.out, square(st.const_tensor_ref(n.x), s)); @@ -2236,6 +2241,9 @@ class Interpreter { case OpCode::CEIL: ops::exec_ceil(std::get(instr.node), st, s); break; + case OpCode::TRUNC: + ops::exec_trunc(std::get(instr.node), st, s); + break; case OpCode::SQUARE: ops::exec_square(std::get(instr.node), st, s); break; diff --git a/backends/mlx/serialization/schema.fbs b/backends/mlx/serialization/schema.fbs index 87ab9aec10a..57195c19b7b 100644 --- a/backends/mlx/serialization/schema.fbs +++ b/backends/mlx/serialization/schema.fbs @@ -845,6 +845,11 @@ table NegNode { out: Tid (required); } +table TruncNode { + x: Tid (required); + out: Tid (required); +} + // ============================================================================= // Math ops - Binary element-wise // ============================================================================= @@ -1191,7 +1196,8 @@ union OpNode { BitwiseXorNode, IfNode, RandomBitsNode, - UpdateAndAttendNode + UpdateAndAttendNode, + TruncNode, // BC: Add new op nodes here (append only) } diff --git a/backends/mlx/test/test_ops.py b/backends/mlx/test/test_ops.py index 91d6f9b130b..6e3e82cf4ee 100644 --- a/backends/mlx/test/test_ops.py +++ b/backends/mlx/test/test_ops.py @@ -4819,6 +4819,7 @@ def create_model(self) -> nn.Module: _UNARY_OP_TESTS = [ {"op_name": "floor", "op_fn": torch.floor, "shapes": _SHAPES_3, "input_fn": _input_fn(scale=10)}, {"op_name": "ceil", "op_fn": torch.ceil, "shapes": _SHAPES_3, "input_fn": _input_fn(scale=10)}, + {"op_name": "trunc", "op_fn": torch.trunc, "shapes": _SHAPES_3, "input_fn": _input_fn(scale=10)}, {"op_name": "square", "op_fn": torch.square, "shapes": _SHAPES_3}, {"op_name": "exp", "op_fn": torch.exp, "shapes": _SHAPES_3}, {"op_name": "sin", "op_fn": torch.sin, "shapes": _SHAPES_3, "input_fn": _input_fn(scale=3.14159)}, From 9030882b28ad8f4fd111aad3bd8afcd9bf7128ee Mon Sep 17 00:00:00 2001 From: Bharath Vedartham Date: Tue, 25 Aug 2026 23:09:56 +0530 Subject: [PATCH 2/2] fix lint --- backends/mlx/ops.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/backends/mlx/ops.py b/backends/mlx/ops.py index 010b1b35b85..ea5f1ac36cc 100644 --- a/backends/mlx/ops.py +++ b/backends/mlx/ops.py @@ -838,6 +838,7 @@ def handler(P: MLXProgramBuilder, n: Node) -> Slot: _make_binary_handler(_node_cls, _op_name, _lift_b) ) + _SCALAR_INT_OPS: List[Tuple[Any, Any, str]] = [ (operator.add, AddIntNode, "operator.add"), (operator.sub, SubtractIntNode, "operator.sub"), @@ -938,6 +939,7 @@ def handler(P: MLXProgramBuilder, n: Node) -> Slot: _make_reduction_handler(_node_cls, _op_name, _max_args) ) + _FULL_OPS: List[Tuple[List[Any], str, Optional[float]]] = [ ([torch.ops.aten.full.default], "aten.full", None), ([torch.ops.aten.zeros.default], "aten.zeros", 0.0), @@ -988,6 +990,7 @@ def handler(P: MLXProgramBuilder, n: Node) -> Slot: for _targets, _op_name, _fixed_fill in _FULL_OPS: REGISTRY.register(target=_targets)(_make_full_handler(_op_name, _fixed_fill)) + _FULL_LIKE_OPS: List[Tuple[List[Any], str, Optional[float]]] = [ ([torch.ops.aten.full_like.default], "aten.full_like", None), ([torch.ops.aten.zeros_like.default], "aten.zeros_like", 0.0),