onnxscript 0.7.1.dev20260422__tar.gz → 0.7.1.dev20260428__tar.gz

{onnxscript-0.7.1.dev20260422/onnxscript.egg-info → onnxscript-0.7.1.dev20260428}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnxscript
-Version: 0.7.1.dev20260422
+Version: 0.7.1.dev20260428
 Summary: Naturally author ONNX functions and models using a subset of Python
 Author-email: Microsoft Corporation <onnx@microsoft.com>
 License: MIT License
@@ -27,7 +27,7 @@ License: MIT License
         
 Project-URL: Homepage, https://microsoft.github.io/onnxscript/
 Project-URL: Repository, https://github.com/microsoft/onnxscript
-Project-URL: Commit, https://github.com/microsoft/onnxscript/tree/1911741522aede0e2eb8269c2e94373b048d32e0
+Project-URL: Commit, https://github.com/microsoft/onnxscript/tree/2277e036d50f1a2bd8351218b19d60315d5a9bcf
 Classifier: Development Status :: 4 - Beta
 Classifier: Environment :: Console
 Classifier: Intended Audience :: Developers

{onnxscript-0.7.1.dev20260422 → onnxscript-0.7.1.dev20260428}/onnxscript/_internal/converter.py

@@ -1478,7 +1478,6 @@ class Converter:
             self._current_fn = irbuilder.IRFunction(stmt.name, domain)
             self._analyzer = analysis.AstAnalyzer(stmt, self._message, self.globals)
             fn_ir = self._translate_function_def_common(stmt)
-            self.this_module.add_function_def(fn_ir)
             self._analyzer = None
             return fn_ir
         raise ValueError(f"Unsupported top-level statement type {type(stmt)!r}.")

{onnxscript-0.7.1.dev20260422 → onnxscript-0.7.1.dev20260428}/onnxscript/_internal/values.py

@@ -94,7 +94,6 @@ class Opset:
         instance = super().__new__(cls)
         instance.domain = domain  # type: ignore[attr-defined]
         instance.version = version  # type: ignore[attr-defined]
-        instance.function_defs = {}  # type: ignore[attr-defined]
         cls.cache[key] = instance
         return instance
 
@@ -130,9 +129,6 @@ class Opset:
         except Exception as exc:  # pylint: disable=broad-exception-caught
             raise AttributeError(f"Attribute {attr} not found.") from exc
 
-    def add_function_def(self, fun):
-        self.function_defs[fun.name] = fun
-
     def _prepare_inputs(self, _: onnx.defs.OpSchema, *inputs):
         """Trims 'None' values from the end of the inputs list. This is used to support
         omitting optional inputs when no more required inputs follow to prepare a valid call

{onnxscript-0.7.1.dev20260422 → onnxscript-0.7.1.dev20260428}/onnxscript/rewriter/models/_rotary_embedding_models.py

@@ -168,3 +168,42 @@ class _PartialRotaryTestCase:
 
 def partial_rotary_test_case():
     return _PartialRotaryTestCase()
+
+
+def _make_partial_rotary_script(mismatched: bool = False):
+    """Generate a partial rotary embedding script with matching or mismatched slice boundaries.
+
+    Args:
+        mismatched: If True, the second slice starts at 33 instead of 32,
+            creating a gap that should prevent PartialRotaryEmbedding fusion.
+    """
+
+    @script()
+    def partial_rotary(position_ids, query):
+        inv_freqs = op.Constant(value=inv_freqs_value)  # [1, rd/2, 1]
+        position_ids_3d = op.Unsqueeze(position_ids, 1)  # [B, 1, S]
+        position_ids_3d_float = op.Cast(position_ids_3d, to=1)
+        matmul = op.MatMul(inv_freqs, position_ids_3d_float)  # [B, rd/2, S]
+        transpose = op.Transpose(matmul, perm=[0, 2, 1])  # [B, S, rd/2]
+        cat = op.Concat(transpose, transpose, axis=-1)  # [B, S, rd]
+        cos_3d = op.Cos(cat)  # [B, S, rd]
+        sin_3d = op.Sin(cat)  # [B, S, rd]
+        # Split the query for partial embedding
+        to_embed = op.Slice(query, [0], [32], [3], [1])
+        if mismatched:
+            unembedded = op.Slice(query, [33], [9223372036854775807], [3], [1])
+        else:
+            unembedded = op.Slice(query, [32], [9223372036854775807], [3], [1])
+        cos_4d = op.Unsqueeze(cos_3d, [1])  # [B, 1, S, rd]
+        sin_4d = op.Unsqueeze(sin_3d, [1])  # [B, 1, S, rd]
+        to_embed_times_cos = op.Mul(to_embed, cos_4d)
+        to_embed_x = op.Slice(to_embed, [0], [16], [3], [1])
+        to_embed_y = op.Slice(to_embed, [16], [9223372036854775807], [3], [1])
+        minus_to_embed_y = op.Neg(to_embed_y)
+        to_embed_rotated_90 = op.Concat(minus_to_embed_y, to_embed_x, axis=-1)
+        to_embed_rotated_90_times_sin = op.Mul(to_embed_rotated_90, sin_4d)
+        embedded = op.Add(to_embed_times_cos, to_embed_rotated_90_times_sin)
+        final = op.Concat(embedded, unembedded, axis=-1)
+        return final
+
+    return partial_rotary

{onnxscript-0.7.1.dev20260422 → onnxscript-0.7.1.dev20260428/onnxscript.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnxscript
-Version: 0.7.1.dev20260422
+Version: 0.7.1.dev20260428
 Summary: Naturally author ONNX functions and models using a subset of Python
 Author-email: Microsoft Corporation <onnx@microsoft.com>
 License: MIT License
@@ -27,7 +27,7 @@ License: MIT License
         
 Project-URL: Homepage, https://microsoft.github.io/onnxscript/
 Project-URL: Repository, https://github.com/microsoft/onnxscript
-Project-URL: Commit, https://github.com/microsoft/onnxscript/tree/1911741522aede0e2eb8269c2e94373b048d32e0
+Project-URL: Commit, https://github.com/microsoft/onnxscript/tree/2277e036d50f1a2bd8351218b19d60315d5a9bcf
 Classifier: Development Status :: 4 - Beta
 Classifier: Environment :: Console
 Classifier: Intended Audience :: Developers