From de165c3ba55a40c4eaeaf7137e0bf0c59b413f49 Mon Sep 17 00:00:00 2001
From: jinchen62 <49575973+jinchen62@users.noreply.github.com>
Date: Tue, 28 Nov 2023 03:25:19 -0800
Subject: [PATCH] Fix (export): update torch-mlir jit_ir import path (#771)

---
 src/brevitas_examples/llm/llm_quant/mlir_custom_mm.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/src/brevitas_examples/llm/llm_quant/mlir_custom_mm.py b/src/brevitas_examples/llm/llm_quant/mlir_custom_mm.py
index 639d5300c..e5ece7b6a 100644
--- a/src/brevitas_examples/llm/llm_quant/mlir_custom_mm.py
+++ b/src/brevitas_examples/llm/llm_quant/mlir_custom_mm.py
@@ -9,10 +9,9 @@
 import torch
 import torch.utils.cpp_extension
 import torch_mlir
-from torch_mlir.dialects.torch.importer.jit_ir.build_tools.registry import \
-    _rename_python_keyword_parameter_name
-from torch_mlir.dialects.torch.importer.jit_ir.build_tools.registry import JitOperator
-from torch_mlir.dialects.torch.importer.jit_ir.build_tools.registry import SIG_ATTR_TYPE
+from torch_mlir.jit_ir_importer.build_tools.registry import _rename_python_keyword_parameter_name
+from torch_mlir.jit_ir_importer.build_tools.registry import JitOperator
+from torch_mlir.jit_ir_importer.build_tools.registry import SIG_ATTR_TYPE
 
 from brevitas.backport.fx._symbolic_trace import wrap