与机器学习框架的整合

import numpy as np
import tvm
from tvm import relax
from tvm.ir.module import IRModule
from tvm.script import relax as R
from tvm.script import tir as T

import torch
import torch.nn as nn
from torch import fx
from torch.nn import functional as F

通过 Builder 构造 IRModule

从张量表达式构造 TensorIR

from tvm import te

A = te.placeholder((128, 128), name="A", dtype="float32")
B = te.placeholder((128, 128), name="B", dtype="float32")
type(A)

tvm.te.tensor.Tensor

每个 te.Tensor 都有 shape 字段和 dtype 字段，用于记录计算的 shape 和数据类型。