def test_convert_follow_node_with_integer_arguments():
"""Tests the conversion of a follow op with integer arguments + constant float args.
The follow op should convert the floating point argument into fp16 as constants/vars
will always be converted if safe to do so.
"""
data = relay.var("data", shape=[1, 10], dtype="float32")
# We use an addition to make sure the input indices are not a var
# (which are always casted if safe)
indices = relay.var("indices", shape=[1, 1], dtype="int32") + relay.const(
0, dtype="int32")
take = relay.take(data, indices, axis=0)
mod = tvm.IRModule.from_expr(take)
mod_params = {
"data": np.random.uniform(-1, 1, size=[1, 10]).astype("float32"),
"indices": np.array([[0]]).astype("int32"),
}
output_mod = verify_mixed_precision_output_close(mod,
mod_params,
atol=0.01,
rtol=0.01)
# Create expected module
data = relay.cast(relay.var("data", shape=[1, 10]), "float16")
take = relay.take(data, indices, axis=0)
expected_mod = tvm.IRModule.from_expr(take)
expected_mod = InferType()(expected_mod)
assert tvm.ir.structural_equal(expected_mod, output_mod)
def test_take_grad(executor_kind):
data_dtype = relay.TensorType((3, 4, 5), "float64")
data = relay.var("data", data_dtype)
indices = relay.var("indices", relay.TensorType((relay.Any(), ), "int32"))
inputs = [
_np_randn_from_type(data_dtype, scale=1e-5),
np.array([1, 2], dtype="int32")
]
test_inputs = [inputs[0]]
# take on axis
fwd_func = relay.Function([data, indices], relay.take(data,
indices,
axis=1))
check_grad(fwd_func,
inputs=inputs,
test_inputs=test_inputs,
executor_kind=executor_kind)
# take on flattened
fwd_func = relay.Function([data, indices],
relay.take(data, indices, axis=None))
check_grad(fwd_func,
inputs=inputs,
test_inputs=test_inputs,
executor_kind=executor_kind)
def verify_take(dshape, indices_shape, oshape, axis=None):
x = relay.var("x", relay.TensorType(dshape, "float32"))
indices = relay.var("indices", relay.TensorType(indices_shape, "int32"))
y = relay.take(x, indices, axis=axis)
y.astext()
yy = relay.ir_pass.infer_type(y)
assert yy.checked_type == relay.TensorType(oshape, "float32")
def verify_take(dshape, indices_shape, oshape, axis=None):
x = relay.var("x", relay.TensorType(dshape, "float32"))
indices = relay.var("indices", relay.TensorType(indices_shape, "int32"))
y = relay.take(x, indices, axis=axis)
y.astext()
yy = relay.ir_pass.infer_type(y)
assert yy.checked_type == relay.TensorType(oshape, "float32")
def test_take(self):
data = relay.var("data", relay.TensorType((-1, 3, 224, 224),
"float32"))
indices = relay.var("indices", relay.TensorType([], "int32"))
net = relay.take(data, indices, axis=1)
net = relay.Function(relay.analysis.free_vars(net), net)
mod = tvm.IRModule.from_expr(net)
mod = relay.transform.InferType()(mod)
xgraph = xf_relay.from_relay(mod, {"indices": np.array(0, np.int32)})
layers = xgraph.get_layers()
assert layers[0].type[0] == "Input"
assert layers[1].type[0] == "Constant"
assert layers[1].data == np.array(0, np.int32)
assert layers[2].type[0] == "Take"
assert "relay_id" in layers[2].attrs
assert layers[2].attrs["axis"] == 1
assert layers[2].attrs["mode"] == "clip"
assert layers[2].shapes == [-1, 224, 224]
def test_arange_with_dynamic_shape():
m, n, k = relay.ShapeVar('m'), relay.ShapeVar('n'), relay.ShapeVar('k')
x = relay.var('x', shape=(m.var, n.var, k.var), dtype='float32')
y0 = relay.shape_of(x)
y1 = relay.take(y0, relay.const(0, 'int32'))
y2 = relay.op.arange(y1)
ex = relay.create_executor()
f = relay.Function([x], y2, type_params=[m, n, k])
def create_integer_lookup_op(
input_arg: "relay.Expr",
floating_point_func: Callable[[np.array], np.array],
in_scale: "relay.Expr",
in_zero_point: "relay.Expr",
out_scale: "relay.Expr",
out_zero_point: "relay.Expr",
in_axis: int = -1,
out_axis: int = -1,
in_dtype: str = "uint8",
out_dtype: str = "uint8",
) -> "relay.Expr":
"""
Create a quantized version of the given floating point unary operation using table lookup.
Args:
input_arg: The quantized input to the final function.
floating_point_func: The numpy function which this table is to approximate
in_scale: The scale of the quantized input tensor.
in_zero_point: The zero point of the quantized input tensor.
out_scale: The scale of the quantized output tensor.
out_zero_point: The zero point of the quantized output tensor.
in_axis: The axis for multi-channel quantization of the input if applicable.
out_axis: The axis for multi-channel quantization of the output if applicable.
in_dtype: The dtype of the input tensor.
out_dtype: The wanted dtype of the output tensor.
Returns:
A Relay expression representing a quantized version of the given function.
"""
# TODO: handle multi-channel q, below will fail with multi-channel q
in_scale = in_scale.data.numpy().item()
in_zero_point = in_zero_point.data.numpy().item()
out_scale = out_scale.data.numpy().item()
out_zero_point = out_zero_point.data.numpy().item()
lookup_table = create_integer_lookup_table(
floating_point_func,
relay.const(in_scale),
relay.const(in_zero_point, dtype="int32"),
relay.const(out_scale),
relay.const(out_zero_point, dtype="int32"),
in_axis=in_axis,
in_dtype=in_dtype,
out_axis=out_axis,
out_dtype=out_dtype,
)
in_dtype_info = np.iinfo(in_dtype)
in_dtype_num_bits = in_dtype_info.bits
lookup_table = relay.const(lookup_table)
index_tensor = relay.reinterpret(input_arg, f"uint{in_dtype_num_bits}")
result = relay.take(lookup_table, index_tensor, axis=0, mode="fast")
return result
def verify_take(dshape, indices_shape, oshape, axis=None):
ib = relay.ir_builder.IRBuilder()
x = ib.param("x", relay.ty.TensorType(dshape, "float32"))
indices = ib.param("indices", relay.ty.TensorType(indices_shape, "int32"))
with ib.function(x, indices) as func:
ib.ret(relay.take(x.var, indices.var, axis=axis))
ib.ret(func)
func = relay.ir_pass.infer_type(ib.env, func.to_func())
ftype = func.checked_type
assert ftype.ret_type == relay.ty.TensorType(oshape, "float32")
def test_arange_with_dynamic_shape():
# m, n, k = relay.ShapeVar('m'), relay.ShapeVar('n'), relay.ShapeVar('k')
m, n, k = relay.Any(), relay.Any(), relay.Any()
x = relay.var("x", shape=(m, n, k), dtype="float32")
y0 = relay.shape_of(x)
y1 = relay.take(y0, relay.const(0, "int32"))
y2 = relay.op.arange(y1, dtype="int32")
y3 = y2 + relay.const(1, dtype="int32")
data = np.random.rand(10, 5, 3).astype("float32")
mod = tvm.IRModule()
mod["main"] = relay.Function([x], y3)
check_result([data], mod, np.array(range(10)).astype("int32") + 1)
def test_let_bound_constants():
"""This tests for an ICHECK failure for ill-formed IR with let-bound constants"""
x = relay.var("x", shape=(3,), dtype="int32")
y = relay.take(x, relay.const(0))
z = relay.const(1)
f = relay.Function([x], relay.stack((z, y), axis=0))
mod = IRModule.from_expr(f)
compiler = VMCompiler()
compiler.optimize(mod, target="llvm")
def verify_any_take(data_shape, indices_shape, axis, data_np_shape, indices_np_shape):
mod = tvm.IRModule()
data = relay.var("data", shape=data_shape, dtype="float32")
indices = relay.var("indices", shape=indices_shape, dtype="int32")
y = relay.take(data, indices, axis=axis)
mod["main"] = relay.Function([data, indices], y)
data_np = np.random.uniform(size=data_np_shape).astype("float32")
if axis is None:
max_index = data_np.size
else:
max_index = data_np.shape[axis]
indices_np = np.random.randint(max_index, size=indices_np_shape).astype("int32")
ref = np.take(data_np, indices_np, axis=axis)
check_result([data_np, indices_np], mod, ref)
def test_arange_with_dynamic_shape():
m, n, k = relay.ShapeVar('m'), relay.ShapeVar('n'), relay.ShapeVar('k')
x = relay.var('x', shape=(m.var, n.var, k.var), dtype='float32')
y0 = relay.shape_of(x)
y1 = relay.take(y0, relay.const(0, 'int32'))
y2 = relay.op.arange(y1, dtype="int32")
y3 = y2 + relay.const(1, dtype="int32")
data = np.random.rand(10, 5, 3).astype('float32')
mod = relay.module.Module()
mod["main"] = relay.Function([x], y3, type_params=[m, n, k])
for kind in ["debug", "vm"]:
ex = relay.create_executor(kind, mod=mod, ctx=tvm.cpu(), target="llvm")
result = ex.evaluate()(data)
tvm.testing.assert_allclose(result.asnumpy(),
np.array(range(10)).astype("int32") + 1)
def test_any_shape_of():
x = relay.var("x", shape=any_dims(2), dtype="float32")
y = relay.shape_of(x)
mod = tvm.IRModule()
mod["main"] = relay.Function([x], y)
data = np.random.uniform(size=(3, 4)).astype("float32")
check_result([data], mod, np.array([3, 4]).astype("int64"))
x = relay.var("x", shape=any_dims(3), dtype="float32")
y0 = relay.shape_of(x)
y1 = relay.take(y0, relay.const(1, "int32"))
mod = tvm.IRModule()
mod["main"] = relay.Function([x], y1)
data = np.random.uniform(size=(2, 3, 4)).astype("float32")
check_result([data], mod, np.array(3).astype("int64"))
def verify_take(src_shape, indices_src, axis=None, mode='clip'):
src_dtype = 'float32'
indices_dtype = 'int32'
indices_src = np.array(indices_src, dtype=indices_dtype)
x = relay.var('x', relay.TensorType(src_shape, src_dtype))
indices = relay.var('indices', relay.TensorType(indices_src.shape, indices_dtype))
z = relay.take(x, indices, axis=axis, mode=mode)
func = relay.Function([x, indices], z)
x_data = np.random.uniform(low=(- 1), high=1, size=src_shape).astype(src_dtype)
np_mode = ('raise' if (mode == 'fast') else mode)
ref_res = np.take(x_data, indices=indices_src, axis=axis, mode=np_mode)
for (target, ctx) in tvm.testing.enabled_targets():
for kind in ['graph', 'debug']:
intrp = relay.create_executor(kind, ctx=ctx, target=target)
op_res = intrp.evaluate(func)(x_data, indices_src)
tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-05)
def verify_take(src_shape, indices_src, axis=None):
src_dtype = "float32"
indices_dtype = "int32"
indices_src = np.array(indices_src, dtype=indices_dtype)
x = relay.var("x", relay.TensorType(src_shape, src_dtype))
indices = relay.var("indices", relay.TensorType(indices_src.shape, indices_dtype))
z = relay.take(x, indices, axis=axis)
func = relay.Function([x, indices], z)
x_data = np.random.uniform(low=-1, high=1, size=src_shape).astype(src_dtype)
ref_res = np.take(x_data, indices=indices_src, axis=axis)
for target, ctx in ctx_list():
for kind in ["graph", "debug"]:
intrp = relay.create_executor(kind, ctx=ctx, target=target)
op_res = intrp.evaluate(func)(x_data, indices_src)
tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
def verify_take(src_shape, indices_src, axis=None):
src_dtype = "float32"
indices_dtype = "int32"
indices_src = np.array(indices_src, dtype=indices_dtype)
x = relay.var("x", relay.TensorType(src_shape, src_dtype))
indices = relay.var("indices", relay.TensorType(indices_src.shape, indices_dtype))
z = relay.take(x, indices, axis=axis)
func = relay.Function([x, indices], z)
x_data = np.random.uniform(low=-1, high=1, size=src_shape).astype(src_dtype)
ref_res = np.take(x_data, indices=indices_src, axis=axis)
for target, ctx in ctx_list():
for kind in ["graph", "debug"]:
intrp = relay.create_executor(kind, ctx=ctx, target=target)
op_res = intrp.evaluate(func)(x_data, indices_src)
tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
def verify_any_take(data_shape, indices_shape, axis, data_np_shape, indices_np_shape):
mod = tvm.IRModule()
data = relay.var('data', shape=data_shape, dtype='float32')
indices = relay.var('indices', shape=indices_shape, dtype='int32')
y = relay.take(data, indices, axis=axis)
mod["main"] = relay.Function([data, indices], y)
data_np = np.random.uniform(size=data_np_shape).astype('float32')
if axis is None:
max_index = data_np.size
else:
max_index = data_np.shape[axis]
indices_np = np.random.randint(max_index, size=indices_np_shape).astype('int32')
ref = np.take(data_np, indices_np, axis=axis)
for kind in ["debug", "vm"]:
ex = relay.create_executor(kind, mod=mod, ctx=tvm.cpu(), target="llvm")
result = ex.evaluate()(data_np, indices_np)
tvm.testing.assert_allclose(result.asnumpy(), ref)
def test_any_shape_of():
x = relay.var('x', shape=any_dims(2), dtype='float32')
y = relay.shape_of(x)
mod = tvm.IRModule()
mod["main"] = relay.Function([x], y)
data = np.random.uniform(size=(3, 4)).astype('float32')
for kind in ["debug", "vm"]:
ex = relay.create_executor(kind, mod=mod, ctx=tvm.cpu(), target="llvm")
result = ex.evaluate()(data)
tvm.testing.assert_allclose(result.asnumpy(), np.array([3,4]).astype("int64"))
x = relay.var('x', shape=any_dims(3), dtype='float32')
y0 = relay.shape_of(x)
y1 = relay.take(y0, relay.const(1, 'int32'))
mod = tvm.IRModule()
mod["main"] = relay.Function([x], y1)
data = np.random.uniform(size=(2, 3, 4)).astype('float32')
for kind in ["debug", "vm"]:
ex = relay.create_executor(kind, mod=mod, ctx=tvm.cpu(), target="llvm")
result = ex.evaluate()(data)
tvm.testing.assert_allclose(result.asnumpy(), np.array(3).astype("int64"))
P.scalar_ge: relay.op.greater_equal,
P.bool_and: relay.op.logical_and,
P.scalar_bit_and: relay.op.bitwise_and,
P.scalar_bit_or: relay.op.bitwise_or,
P.scalar_bit_xor: relay.op.bitwise_xor,
P.scalar_bit_not: relay.op.bitwise_not,
P.scalar_bit_lshift: relay.op.left_shift,
P.scalar_bit_rshift: relay.op.right_shift,
P.bool_or: relay.op.logical_or,
P.bool_eq: relay.op.equal,
P.bool_not: relay.op.logical_not,
P.array_to_scalar: lambda x: x,
P.dot: lambda x, y: relay.op.nn.dense(x, relay.op.transpose(y)),
P.make_tuple: lambda *args: relay.Tuple(args),
P.switch: relay.If,
P.take: lambda w, i: relay.take(w, i, 0),
}
def relay_distribute(c, array, shape):
"""Implementation of distribute for Relay."""
assert shape.is_constant(tuple)
# Make sure shape is a tuple of builtin Python integers.
relay_shape = tuple(int(dim) for dim in shape.value)
return relay.op.broadcast_to(c.ref(array), relay_shape)
def relay_transpose(c, a, ax):
"""Implementation of transpose for Relay."""
na = c.ref(a)
assert ax.is_constant(tuple)
def flexible_dispatch(mod,
buckets,
axis=0,
auto_pad=False,
pad_value=0,
input_indices=None,
affects_output=True):
"""
Enable inference of multiple shaped inputs in one module.
This transformation adds a handler around a module that
checks input shapes and dispatches to a subgraph specialized
to handle the specific shapes of that input. If no exactly matching
subgraph is available, the input will be run using full dynamism.
For best performance, specify all the sizes the module will
be likely to see using the buckets argument.
By default, this function will dispatch shapes that exactly match one
of the buckets to a corresponding subgraph. All non-matching shapes
use the same fully dynamic fallback. This can be detrimental to performance
for those non-matching shapes. Setting auto_pad to True causes this
function to round-up the shape of non-matching inputs to the closest
bucket. This allows them to use the tuned kernels of bucket shapes
which can improve performance.
Functions that have multiple inputs sharing a dynamic axis, which
is common for batch size or sequence length dynamism, are supported
through the input_indices argument.
Many types of dynamism such as batching affect both the input and output
shape, however this is not always the case. If the output shape
is independent of the input, the affects_output argument of this
function must be set to False.
Parameters
----------
buckets: list[int]
The sizes of the input dimension that should be explicitly handled.
Each value in buckets will have a corresponding subgraph constructed to
handle it.
axis: int
The dimension of the input that should be made flexible. This will
most often be used for the batch dimension.
auto_pad: Optional[bool]
If True, then padding will be inserted to values that don't match one of
the provided buckets.
pad_value: Optional[float]
When auto_pad is true, padding will be done with this value.
input_indices: Optional[List[int]]
Which inputs should be dispatched dynamically, provided by index. All inputs
must share the same dynamic axis.
affects_output: Optional[bool]
Whether the change in input shape has a corresponding effect on the output shape.
Batching for example effects both the input and output whereas changing sequence
length in an NLP model typically does not.
Returns
-------
mod : IRModule
The new module wrapped with a flexible shape dispatch handler.
"""
main_fn = mod["main"]
# Default to single input if not specified.
if input_indices is None:
input_indices = [0]
# Extract all input data and create a new dynamic variable for each.
data = []
dyn_data = []
for i in input_indices:
data.append(main_fn.params[i])
dyn_shape = override_shape(data[i].type_annotation, axis, relay.Any())
dyn_data.append(relay.Var(data[i].name_hint,
type_annotation=dyn_shape))
# Extract the dynamic shape value from one of the inputs.
rt_sh = relay.op.shape_of(dyn_data[0])
flex_value = relay.op.take(rt_sh, relay.const(axis))
if_exprs = []
for i, bucket in enumerate(buckets):
input_data = dyn_data
check_dim = flex_value
# Apply automatic padding if specified.
if auto_pad:
input_data = []
# Construct padding expression for inputs.
for j, inp in enumerate(dyn_data):
pad_width = relay.const(bucket) - flex_value
rank = len(data[j].type_annotation.shape)
pads = relay.zeros([rank, 2], "int32")
pads = relay.scatter_nd(pads, relay.const([axis, 1]),
pad_width)
padded_value = relay.nn.pad(inp, pads, pad_value)
# Determine if this is the proper bucket to pad to. Do this by checking if the
# input shape is between this bucket and the previous.
if i == 0:
padded_value = relay.If(
relay.op.less_equal(flex_value, relay.const(bucket)),
padded_value, inp)
else:
padded_value = relay.If(
relay.op.logical_and(
relay.op.less_equal(flex_value,
relay.const(bucket)),
relay.op.greater(flex_value,
relay.const(buckets[i - 1])),
),
padded_value,
inp,
)
# Update input value and test dimension to reflect possible padding.
input_data.append(padded_value)
# Grab the new possibly padded shape for checking bucket size.
check_dim = relay.op.take(relay.op.shape_of(input_data[0]),
relay.const(axis))
# Create a specialized subgraph for the current bucket.
spec_call, spec_ty = specialize_body(mod,
main_fn,
axis,
bucket,
input_indices=input_indices,
affects_output=affects_output)
# Apply hard casting to shape to create statically typed graphs.
spec_data = []
for j, inp in enumerate(input_data):
spec_data.append(relay.op.reshape(inp, spec_ty[j].shape))
# Create a dispatch statement for the current specialized graph.
call_args = list(main_fn.params)
for j, inp in enumerate(input_indices):
call_args[inp] = spec_data[j]
new_call = spec_call(*call_args)
# Remove meaningless padded outputs if applicable.
if auto_pad and affects_output:
new_call = relay.take(
new_call,
relay.arange(start=relay.const(0),
stop=flex_value,
dtype="int32"),
axis=axis,
)
# Add this new case to the dispatch handler.
if_exprs.append((relay.op.equal(check_dim,
relay.const(bucket)), new_call))
# Create a subgraph to handle all other shapes.
default_dyn_call, _ = specialize_body(mod,
main_fn,
axis,
relay.Any(),
input_indices=input_indices,
affects_output=affects_output)
call_args = list(main_fn.params)
for j, inp in enumerate(input_indices):
call_args[inp] = dyn_data[j]
new_body = default_dyn_call(*call_args)
# Create an If chain to dispatch shapes to the appropriate specialized subgraph.
for cond, true_branch in if_exprs:
new_body = relay.If(cond, true_branch, new_body)
# Assign new parameters to the function.
new_params = list(main_fn.params)
for j, inp in enumerate(input_indices):
new_params[inp] = dyn_data[j]
# Update the output shape to be dynamic if needed.
if affects_output:
dyn_ret_type = override_shape(main_fn.ret_type, axis, relay.Any())
else:
dyn_ret_type = main_fn.ret_type
# Assign the handler as the new entrypoint in the module.
new_main = relay.Function(new_params, new_body, dyn_ret_type,
main_fn.type_params, main_fn.attrs)
mod["main"] = new_main
# Do type inference to make sure everything worked.
mod = relay.transform.InferType()(mod)
return mod
def verify_take(dshape, indices_shape, oshape, axis=None):
x = relay.var('x', relay.TensorType(dshape, 'float32'))
indices = relay.var('indices', relay.TensorType(indices_shape, 'int32'))
y = relay.take(x, indices, axis=axis)
yy = run_infer_type(y)
assert (yy.checked_type == relay.TensorType(oshape, 'float32'))
def _execute(self):
self.node_dict = {}
# self.node_dict['1'] = relay.const(np.zeros((1, 128)), dtype='int32')
gelu_a = relay.var('gelu_a', shape=())
gelu_b = relay.var('gelu_b', shape=())
gelu_c = relay.var('gelu_c', shape=())
gelu_d = relay.var('gelu_d', shape=())
gelu_e = relay.var('gelu_e', shape=())
self.node_dict['1'] = relay.var('input.1', shape=(1,128), dtype='int32')
self.node_dict['2'] = relay.var('input.2', shape=(1,128), dtype='int32')
for gnode in self.graph:
name = gnode['name']
op_type = gnode['op_type']
attrs = gnode['attrs']
del attrs['A_shape']
del attrs['O_shape']
inputs = gnode['inputs']
if op_type == 'Const':
arr = np.zeros(attrs['shape'], dtype=np.int32)
y = relay.const(arr, dtype='int32')
elif op_type == 'expand_dims':
x = get_input(self.node_dict, self.params, inputs[0])
y = relay.expand_dims(x, attrs['axis'], attrs['num_newaxis'])
elif op_type == 'reshape':
x = get_input(self.node_dict, self.params, inputs[0])
y = relay.reshape(x, attrs['newshape'])
elif op_type == 'take':
data = get_input(self.node_dict, self.params, inputs[0])
indices = get_input(self.node_dict, self.params, inputs[1])
y = relay.take(data, indices, axis=attrs['axis'][0], mode=attrs['mode'])
elif op_type == 'one_hot':
x = get_input(self.node_dict, self.params, inputs[0])
cc1 = get_input(self.node_dict, self.params, inputs[1])
cc2 = get_input(self.node_dict, self.params, inputs[2])
y = relay.one_hot(x, cc1, cc2, **attrs)
elif op_type == 'strided_slice':
x = get_input(self.node_dict, self.params, inputs[0])
y = relay.strided_slice(x, **attrs)
elif op_type == 'mean':
x = get_input(self.node_dict, self.params, inputs[0])
y = relay.mean(x, axis=attrs['axis'],
exclude=attrs['exclude'],
keepdims=attrs['keepdims'])
elif op_type == 'nn.dense':
x = get_input(self.node_dict, self.params, inputs[0])
weight = get_input(self.node_dict, self.params, inputs[1])
y = relay.nn.dense(x, weight, units=attrs['units'][0])
elif op_type == 'add':
x1 = get_input(self.node_dict, self.params, inputs[0])
x2 = get_input(self.node_dict, self.params, inputs[1])
y = relay.add(x1, x2)
elif op_type == 'subtract':
x1 = get_input(self.node_dict, self.params, inputs[0])
x2 = get_input(self.node_dict, self.params, inputs[1])
y = relay.subtract(x1, x2)
elif op_type == 'multiply':
x1 = get_input(self.node_dict, self.params, inputs[0])
x2 = get_input(self.node_dict, self.params, inputs[1])
y = relay.multiply(x1, x2)
elif op_type == 'power':
x1 = get_input(self.node_dict, self.params, inputs[0])
x2 = get_input(self.node_dict, self.params, inputs[1])
y = relay.power(x1, x2)
elif op_type == 'transpose':
x = get_input(self.node_dict, self.params, inputs[0])
y = relay.transpose(x, **attrs)
elif op_type == 'tanh':
x = get_input(self.node_dict, self.params, inputs[0])
y = relay.tanh(x)
elif op_type == 'squeeze':
x = get_input(self.node_dict, self.params, inputs[0])
y = relay.squeeze(x, **attrs)
elif op_type == 'nn.batch_matmul':
x1 = get_input(self.node_dict, self.params, inputs[0])
x2 = get_input(self.node_dict, self.params, inputs[1])
y = relay.nn.batch_matmul(x1, x2)
elif op_type == 'nn.softmax':
x = get_input(self.node_dict, self.params, inputs[0])
y = relay.nn.softmax(x, **attrs)
elif op_type == 'gelu':
x = get_input(self.node_dict, self.params, inputs[0])
y = x * gelu_a * (gelu_b + relay.tanh(
( gelu_c * (x + gelu_d *
relay.power(x, gelu_e)))))
else:
import pdb; pdb.set_trace()
print( 'not supported op %s ' % op_type)
self.node_dict[name] = y
output_name = self.output_node_ids[0]
output = self.node_dict[output_name]
inputs = relay.analysis.free_vars(output)
# inputs = [self.node_dict['1'], self.node_dict['2']]
func = relay.Function(inputs, output)
mod = tvm.IRModule()
mod['main'] = func
with relay.build_config(opt_level=0):
graph, lib, params = relay.build(mod, 'llvm', params={})
self.m = graph_runtime.create(graph, lib, tvm.cpu())
