def _verify(indices_shape, depth, on_value, off_value, axis, dtype):
indices = relay.var("indices",
relay.TensorType(indices_shape, "int32"))
depth_var = relay.var("depth", relay.TensorType((), "int32"))
on_value_const = relay.const(on_value)
off_value_const = relay.const(off_value)
out = relay.one_hot(indices, on_value_const, off_value_const,
depth_var, axis, dtype)
func = relay.Function([indices, depth_var], out)
indices_np = np.random.randint(0, depth,
size=indices_shape).astype("int32")
out_np = tvm.topi.testing.one_hot(indices_np, on_value, off_value,
depth, axis, dtype)
for target, ctx in tvm.testing.enabled_targets():
if (target != 'cuda'
): #skip cuda because we don't have dynamic support for GPU
for kind in ["vm", "debug"]:
mod = tvm.ir.IRModule.from_expr(func)
intrp = relay.create_executor(kind,
mod=mod,
ctx=ctx,
target=target)
out_relay = intrp.evaluate()(
indices_np, np.array(depth).astype("int32"))
tvm.testing.assert_allclose(out_relay.asnumpy(), out_np)
def _verify(indices_shape, depth, on_value, off_value, axis, dtype):
indices = relay.var("indices",
relay.TensorType(indices_shape, "int32"))
depth_var = relay.const(depth)
on_value_var = relay.var("on_value", relay.TensorType((), "int32"))
off_value_var = relay.var("off_value", relay.TensorType((), "int32"))
out = relay.one_hot(indices, on_value_var, off_value_var, depth_var,
axis, dtype)
params = {
"on_value": on_value,
"off_value": off_value,
}
func = relay.Function([indices, on_value_var, off_value_var], out)
func2 = run_opt_pass(
run_opt_pass(func, transform.DynamicToStatic(), params),
transform.InferType())
zz = func2.body
assert isinstance(zz, relay.Call)
assert zz.op == relay.op.get("one_hot")
indices_np = np.random.randint(0, depth,
size=indices_shape).astype("int32")
out_np = tvm.topi.testing.one_hot(indices_np, on_value, off_value,
depth, axis, dtype)
verify_func(func2, [indices_np], out_np)
def test_one_hot_grad(executor_kind, target, dev, index_dtype, val_dtype):
indices_shape = (3, 4)
depth = 5
axis = -1
inputs = [
np.random.randint(depth, size=indices_shape, dtype=index_dtype),
np.array(np.random.randn() * 1e-5).astype(val_dtype),
np.array(np.random.randn() * 1e-5).astype(val_dtype),
]
test_inputs = inputs[1:]
indices = relay.var("indices", shape=indices_shape, dtype=index_dtype)
on_val = relay.var("on_val", shape=tuple(), dtype=val_dtype)
off_val = relay.var("off_val", shape=tuple(), dtype=val_dtype)
y = relay.one_hot(indices, on_val, off_val, depth, axis, val_dtype)
f = relay.Function([indices, on_val, off_val], y)
check_grad(
f,
inputs=inputs,
test_inputs=test_inputs,
target_devices=[(target, dev)],
executor_kind=executor_kind,
)
def test_one_hot_grad():
indices_shape = (3, 4)
depth = 5
axis = -1
for indices_dtype in ["int32", "int64"]:
for val_dtype in ["float32", "float64"]:
inputs = [
np.random.randint(depth,
size=indices_shape,
dtype=indices_dtype),
np.array(np.random.randn() * 1e-5).astype(val_dtype),
np.array(np.random.randn() * 1e-5).astype(val_dtype),
]
test_inputs = inputs[1:]
indices = relay.var("indices",
shape=indices_shape,
dtype=indices_dtype)
on_val = relay.var("on_val", shape=tuple(), dtype=val_dtype)
off_val = relay.var("off_val", shape=tuple(), dtype=val_dtype)
y = relay.one_hot(indices, on_val, off_val, depth, axis, val_dtype)
f = relay.Function([indices, on_val, off_val], y)
check_grad(f, inputs=inputs, test_inputs=test_inputs)
def _verify(indices_shape, depth, on_value, off_value, axis, dtype):
indices = relay.var("indices", relay.TensorType(indices_shape, "int32"))
on_value_const = relay.const(on_value)
off_value_const = relay.const(off_value)
out = relay.one_hot(indices, on_value_const, off_value_const, depth, axis, dtype)
checked = run_infer_type(out)
assert checked.checked_type == relay.ty.TensorType(_get_oshape(indices_shape, depth, axis), dtype)
func = relay.Function([indices], out)
indices_np = np.random.randint(0, depth, size=indices_shape).astype("int32")
out_np = topi.testing.one_hot(indices_np, on_value, off_value, depth, axis, dtype)
for target, ctx in ctx_list():
for kind in ["graph", "debug"]:
intrp = relay.create_executor(kind, ctx=ctx, target=target)
out_relay = intrp.evaluate(func)(indices_np)
tvm.testing.assert_allclose(out_relay.asnumpy(), out_np)
def _verify(indices_shape, depth, on_value, off_value, axis, dtype):
indices = relay.var("indices",
relay.TensorType(indices_shape, "int32"))
depth_var = relay.var("depth", relay.TensorType((), "int32"))
on_value_const = relay.const(on_value)
off_value_const = relay.const(off_value)
out = relay.one_hot(indices, on_value_const, off_value_const,
depth_var, axis, dtype)
func = relay.Function([indices, depth_var], out)
indices_np = np.random.randint(0, depth,
size=indices_shape).astype("int32")
out_np = tvm.topi.testing.one_hot(indices_np, on_value, off_value,
depth, axis, dtype)
for target, dev in tvm.testing.enabled_targets():
mod = tvm.ir.IRModule.from_expr(func)
out_relay = relay.create_executor(
executor_kind, mod=mod, device=dev,
target=target).evaluate()(indices_np,
np.array(depth).astype("int32"))
tvm.testing.assert_allclose(out_relay.numpy(), out_np)
def _verify(indices_shape, depth, on_value, off_value, axis, dtype):
indices = relay.var("indices",
relay.TensorType(indices_shape, "int32"))
on_value_const = relay.const(on_value)
off_value_const = relay.const(off_value)
out = relay.one_hot(indices, on_value_const, off_value_const, depth,
axis, dtype)
checked = run_infer_type(out)
assert checked.checked_type == relay.ty.TensorType(
_get_oshape(indices_shape, depth, axis), dtype)
func = relay.Function([indices], out)
indices_np = np.random.randint(0, depth,
size=indices_shape).astype("int32")
out_np = tvm.topi.testing.one_hot(indices_np, on_value, off_value,
depth, axis, dtype)
for target, dev in tvm.testing.enabled_targets():
out_relay = relay.create_executor(
executor_kind, device=dev,
target=target).evaluate(func)(indices_np)
tvm.testing.assert_allclose(out_relay.numpy(), out_np)
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())
