gimlet-api 0.0.1__py3-none-any.whl

gimlet_api-0.0.1.dist-info/METADATA

@@ -0,0 +1,6 @@
+Metadata-Version: 2.1
+Name: gimlet-api
+Version: 0.0.1
+Requires-Dist: torch-mlir-gml
+Requires-Dist: torch
+

gimlet_api-0.0.1.dist-info/RECORD

@@ -0,0 +1,6 @@
+gml/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+gml/compile.py,sha256=VzNVLojb8D8lISLYxazR3-YkJ6HzscP0juP7yWAltfg,869
+gimlet_api-0.0.1.dist-info/METADATA,sha256=F7JACoU323q1pnbZD-KQfdJTpd_yKFBS8sXhSHhx-M4,106
+gimlet_api-0.0.1.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+gimlet_api-0.0.1.dist-info/top_level.txt,sha256=wSbIUt3C0p7m2JMxpjK6M83pB9pqGe274eF3qgIxzS8,4
+gimlet_api-0.0.1.dist-info/RECORD,,

gimlet_api-0.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: bdist_wheel (0.43.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

gimlet_api-0.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+gml

gml/compile.py

@@ -0,0 +1,26 @@
+from torch.fx.experimental.proxy_tensor import make_fx
+
+#from torch_mlir.torchscript import compile as torch_mlir_compile
+from torch_mlir import compile as torch_mlir_compile
+from torch_mlir import OutputType
+from torch_mlir.dynamo import _get_decomposition_table
+
+
+def to_torch_mlir(model, example_inputs):
+    # Running the model a few times on the inputs, leads to more consistent compiled results.
+    # TODO(james): would be good to understand why this is the case.
+    for _ in range(2):
+        _ = model(*example_inputs)
+
+    model = make_fx(model, pre_dispatch=True, decomposition_table=_get_decomposition_table())(*example_inputs)
+
+    compiled = torch_mlir_compile(
+        model,
+        example_inputs,
+        use_tracing=False,
+        ignore_traced_shapes=False,
+        output_type=OutputType.RAW,
+        use_make_fx=False,
+    )
+
+    return compiled