Skip to content

Instantly share code, notes, and snippets.

func.func @torch.aten.gather(%arg0: !torch.vtensor<[?,?,?],f32>, %arg1: !torch.vtensor<[?,?,?,?],si64>) -> !torch.vtensor<[?,?,?],f32> {
%int-1 = torch.constant.int -1
%false = torch.constant.bool false
%0 = torch.aten.gather %arg0, %int-1, %arg1, %false : !torch.vtensor<[?,?,?],f32>, !torch.int, !torch.vtensor<[?,?,?,?],si64>, !torch.bool -> !torch.vtensor<[?,?,?],f32>
return %0 : !torch.vtensor<[?,?,?],f32>
}
module attributes {torch.debug_module_name = "_lambda"} {
func.func @forward(%arg0: tensor<1x128xi64>) -> tensor<1x2xf32> {
%0 = "tosa.const"() {value = dense<[[65536, 512, 1]]> : tensor<1x3xi32>} : () -> tensor<1x3xi32>
%1 = "tosa.const"() {value = dense_resource<__elided__> : tensor<2x768xf32>} : () -> tensor<2x768xf32>
%2 = "tosa.const"() {value = dense_resource<__elided__> : tensor<768x768xf32>} : () -> tensor<768x768xf32>
%3 = "tosa.const"() {value = dense_resource<__elided__> : tensor<768xf32>} : () -> tensor<768xf32>
%4 = "tosa.const"() {value = dense_resource<__elided__> : tensor<768x3072xf32>} : () -> tensor<768x3072xf32>
%5 = "tosa.const"() {value = dense_resource<__elided__> : tensor<3072xf32>} : () -> tensor<3072xf32>
%6 = "tosa.const"() {value = dense_resource<__elided__> : tensor<3072x768xf32>} : () -> tensor<3072x768xf32>
%7 = "tosa.const"() {value = dense<-3.40282347E+38> : tensor<f32>} : () -> tensor<f32>
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
from torch.fx.experimental.proxy_tensor import make_fx
from torch._decomp import get_decompositions
import tempfile
import torch_mlir
def prepare_sentence_tokens(hf_model: str, sentence: str):
tokenizer = AutoTokenizer.from_pretrained(hf_model)
#loc = loc(unknown)
module attributes {torch.debug_module_name = "_lambda"} {
func.func private @__torch__.torch.fx.graph_module._lambda.__code_getter(%arg0: !torch.nn.Module<"__torch__.torch.fx.graph_module._lambda"> loc(unknown)) -> !torch.str {
%96 = torch.prim.GetAttr %arg0["_code"] : !torch.nn.Module<"__torch__.torch.fx.graph_module._lambda"> -> !torch.str loc(#loc)
return %96 : !torch.str loc(#loc)
} loc(#loc)
func.func private @__torch__.torch.fx.graph_module._lambda.forward(%arg0: !torch.nn.Module<"__torch__.torch.fx.graph_module._lambda"> loc(unknown), %arg1: !torch.tensor {torch.type_bound = !torch.vtensor<[1,128],si64>} loc(unknown)) -> !torch.tensor {
%true_0 = torch.constant.bool true loc(#loc1)
%int11 = torch.constant.int 11 loc(#loc2)
%int-2 = torch.constant.int -2 loc(#loc3)
#loc = loc(unknown)
module attributes {torch.debug_module_name = "_lambda"} {
func.func private @__torch__.torch.fx.graph_module._lambda.__code_getter(%arg0: !torch.nn.Module<"__torch__.torch.fx.graph_module._lambda"> loc(unknown)) -> !torch.str {
%96 = torch.prim.GetAttr %arg0["_code"] : !torch.nn.Module<"__torch__.torch.fx.graph_module._lambda"> -> !torch.str loc(#loc)
return %96 : !torch.str loc(#loc)
} loc(#loc)
func.func private @__torch__.torch.fx.graph_module._lambda.forward(%arg0: !torch.nn.Module<"__torch__.torch.fx.graph_module._lambda"> loc(unknown), %arg1: !torch.tensor {torch.type_bound = !torch.vtensor<[1,128],si64>} loc(unknown)) -> !torch.tensor {
%true_0 = torch.constant.bool true loc(#loc1)
%int11 = torch.constant.int 11 loc(#loc2)
%int-2 = torch.constant.int -2 loc(#loc3)
#loc = loc(unknown)
module attributes {torch.debug_module_name = "_lambda"} {
func.func @forward(%arg0: !torch.vtensor<[1,128],si64> loc(unknown)) -> !torch.vtensor<[1,2],f32> {
%int1 = torch.constant.int 1 loc(#loc1)
%int0 = torch.constant.int 0 loc(#loc2)
%int-1 = torch.constant.int -1 loc(#loc3)
%true = torch.constant.bool true loc(#loc4)
%none = torch.constant.none loc(#loc)
%false = torch.constant.bool false loc(#loc5)
%int128 = torch.constant.int 128 loc(#loc6)
func.func @forward(%arg0: !torch.vtensor<[1,128],si64>) -> !torch.vtensor<[1,2],f32> {
%0 = builtin.unrealized_conversion_cast %arg0 : !torch.vtensor<[1,128],si64> to tensor<1x128xi64>
%int1 = torch.constant.int 1
%1 = builtin.unrealized_conversion_cast %int1 : !torch.int to i64
%int0 = torch.constant.int 0
%2 = builtin.unrealized_conversion_cast %int0 : !torch.int to i64
%int-1 = torch.constant.int -1
%3 = builtin.unrealized_conversion_cast %int-1 : !torch.int to i64
%true = torch.constant.bool true
%4 = builtin.unrealized_conversion_cast %true : !torch.bool to i1
func.func @torch.aten.as_strided(%arg0: !torch.vtensor<[1,128,2304],f32>) -> !torch.vtensor<[1,128,768],f32> {
%int0 = torch.constant.int 0
%int1 = torch.constant.int 1
%int128 = torch.constant.int 128
%int768 = torch.constant.int 768
%int2304 = torch.constant.int 2304
%240 = torch.prim.ListConstruct %int1, %int128, %int768 : (!torch.int, !torch.int, !torch.int) -> !torch.list<int>
%241 = torch.prim.ListConstruct %int294912, %int2304, %int1 : (!torch.int, !torch.int, !torch.int) -> !torch.list<int>
%242 = torch.aten.as_strided %arg0, %240, %241, %int0 : !torch.vtensor<[1,128,2304],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[1,128,768],f32>
return %242 : !torch.vtensor<[1,128,768],f32>
module attributes {torch.debug_module_name = "_lambda"} {
func.func @forward(%arg0: tensor<1x5xi64>) -> tensor<1x5x50257xf32> {
%0 = "tosa.const"() {value = dense<-3.40282347E+38> : tensor<f32>} : () -> tensor<f32>
%1 = "tosa.const"() {value = dense_resource<__elided__> : tensor<1x1x1024x1024xui8>} : () -> tensor<1x1x1024x1024xi8>
%2 = "tosa.const"() {value = dense<8.000000e+00> : tensor<f32>} : () -> tensor<f32>
%3 = "tosa.const"() {value = dense_resource<__elided__> : tensor<50257x768xf32>} : () -> tensor<50257x768xf32>
%4 = "tosa.const"() {value = dense<7.680000e+02> : tensor<1xf32>} : () -> tensor<1xf32>
%5 = "tosa.const"() {value = dense<[0, 2, 1, 3]> : tensor<4xi64>} : () -> tensor<4xi64>
%6 = "tosa.const"() {value = dense<[0, 1, 3, 2]> : tensor<4xi32>} : () -> tensor<4xi32>
%7 = "tosa.const"() {value = dense<0> : tensor<1x1x5x5xi8>} : () -> tensor<1x1x5x5xi8>
module attributes {torch.debug_module_name = "_lambda"} {
func.func private @__torch__.torch.fx.graph_module._lambda.__code_getter(%arg0: !torch.nn.Module<"__torch__.torch.fx.graph_module._lambda">) -> !torch.str {
%186 = torch.prim.GetAttr %arg0["_code"] : !torch.nn.Module<"__torch__.torch.fx.graph_module._lambda"> -> !torch.str
return %186 : !torch.str
}
func.func private @__torch__.torch.fx.graph_module._lambda.forward(%arg0: !torch.nn.Module<"__torch__.torch.fx.graph_module._lambda">, %arg1: !torch.tensor {torch.type_bound = !torch.vtensor<[1,5],si64>}) -> !torch.tensor {
%int11 = torch.constant.int 11
%int-2 = torch.constant.int -2
%none_0 = torch.constant.none
%false = torch.constant.bool false