chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,948 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
# ruff: noqa: F401, F811, RUF005
|
||||
from collections.abc import Callable
|
||||
from typing import Union
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import tvm_ffi
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm import relax
|
||||
from tvm.ir import Call
|
||||
from tvm.ir.op import Op
|
||||
from tvm.relax.transform import LegalizeOps
|
||||
from tvm.relax.type import TensorType, TupleType
|
||||
from tvm.testing.utils import check_numerical_grads
|
||||
|
||||
|
||||
def relax_check_gradients(
|
||||
op_func: Callable,
|
||||
inputs_numpy: list[np.array],
|
||||
target: str | tvm.target.Target,
|
||||
dev: tvm.runtime.Device,
|
||||
tuple_input: bool = False,
|
||||
ignore_grads: list[int] = [],
|
||||
**kwargs, # attr for operators
|
||||
):
|
||||
"""Generate the forward and the gradient module. Then run them and check numeric gradients.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
op_func : Callable
|
||||
The forward operator function. Should be a function in package relax.op.
|
||||
|
||||
inputs_numpy : List[np.array]
|
||||
The np array inputs for op_func. inputs_numpy will be transformed into TVM Tensor inside
|
||||
this function.
|
||||
|
||||
If op_func takes a tuple of tensors as input, you can set tuple_input as True, and pass the
|
||||
tuple input (or list) as inputs_numpy. See test_concat().
|
||||
|
||||
target : Union[str, tvm.target.Target]
|
||||
The building target.
|
||||
|
||||
dev : tvm.runtime.Device
|
||||
The device to deploy the module.
|
||||
|
||||
tuple_input : bool
|
||||
Whether the operator accepts a tuple as input. If true, operator will accept exactly one
|
||||
tuple of tensors as input; otherwise, operator accept one or more tensors as input. See
|
||||
test_concat(). Default: False.
|
||||
|
||||
ignore_grads: List[int]
|
||||
Specifies which input we do not need to find gradient.
|
||||
|
||||
Sometimes the input is not differentiable, such as shape, boolean values, positions, etc.
|
||||
We can specify the index of these inputs to check the gradient of them is no_grad, and
|
||||
prevent computing numeric gradient.
|
||||
|
||||
kwargs : Any
|
||||
The keyword arguments for the op_func. Will be passed to op_func directly.
|
||||
"""
|
||||
|
||||
func_name = "main"
|
||||
|
||||
# Helper functions
|
||||
def _numpy_to_ty(data):
|
||||
if isinstance(data, list):
|
||||
return relax.TupleType([_numpy_to_ty(d) for d in data])
|
||||
return relax.TensorType(data.shape, str(data.dtype))
|
||||
|
||||
def _numpy_to_tvm(data):
|
||||
if isinstance(data, list):
|
||||
return [_numpy_to_tvm(d) for d in data]
|
||||
return tvm.runtime.tensor(data)
|
||||
|
||||
def _tvm_to_numpy(data, ignore_idx=[]):
|
||||
if isinstance(data, tvm_ffi.Array):
|
||||
return [_tvm_to_numpy(d) for i, d in enumerate(data) if i not in ignore_idx]
|
||||
if isinstance(data, tvm.runtime.Tensor):
|
||||
return data.numpy()
|
||||
return data
|
||||
|
||||
def _gen_weights(out_ty):
|
||||
if isinstance(out_ty, TupleType):
|
||||
return [_gen_weights(ty) for ty in out_ty.fields]
|
||||
else:
|
||||
assert isinstance(out_ty, TensorType)
|
||||
return np.random.uniform(size=[int(i) for i in out_ty.shape]).astype(out_ty.dtype)
|
||||
|
||||
def _is_call_no_grad(expr):
|
||||
return isinstance(expr, Call) and expr.op == Op.get("relax.grad.no_grad")
|
||||
|
||||
# Generate parameter relax Vars
|
||||
param_vars = [
|
||||
relax.Var("x_" + str(i), _numpy_to_ty(data)) for i, data in enumerate(inputs_numpy)
|
||||
]
|
||||
|
||||
# Generate the forward call
|
||||
if tuple_input:
|
||||
t = relax.Tuple(param_vars)
|
||||
call = op_func(t, **kwargs)
|
||||
else:
|
||||
call = op_func(*param_vars, **kwargs)
|
||||
|
||||
# Forward mod
|
||||
forward_bb = relax.BlockBuilder()
|
||||
with forward_bb.function(func_name, param_vars):
|
||||
with forward_bb.dataflow():
|
||||
out = forward_bb.emit_output(call)
|
||||
forward_bb.emit_func_output(out)
|
||||
forward_mod = forward_bb.get()
|
||||
forward_ex = tvm.compile(forward_mod, target)
|
||||
forward_vm = relax.VirtualMachine(forward_ex, dev)
|
||||
|
||||
# Generate weights
|
||||
# In forward process, weights represent the weight of every element of the result of the
|
||||
# forward call. The weighted result will be sum(weight * result).
|
||||
# If the result is a tuple, weights will be a list, and the weighted result will be
|
||||
# sum(i * j for i, j in zip(weights, result))
|
||||
# In the gradient process, weights is the output gradient, i.e. the gradient w.r.t. the result.
|
||||
out_ty = forward_mod[func_name].body.body.ty
|
||||
weights = _gen_weights(out_ty)
|
||||
|
||||
# The inputs of the forward function are inputs_filtered below.
|
||||
def forward(*inputs):
|
||||
inputs_iter = iter(inputs)
|
||||
inputs_tvm = [
|
||||
_numpy_to_tvm(next(inputs_iter))
|
||||
if i not in ignore_grads
|
||||
else _numpy_to_tvm(inputs_numpy[i])
|
||||
for i in range(len(inputs_numpy))
|
||||
]
|
||||
result = forward_vm[func_name](*inputs_tvm)
|
||||
result_numpy = _tvm_to_numpy(result)
|
||||
if isinstance(result_numpy, list):
|
||||
assert isinstance(weights, list)
|
||||
assert len(weights) == len(result_numpy)
|
||||
ret = 0
|
||||
for i, weight in enumerate(weights):
|
||||
ret += np.sum(weight * result_numpy[i])
|
||||
return ret
|
||||
return np.sum(weights * result_numpy)
|
||||
|
||||
# The gradient function
|
||||
assert isinstance(call.op, Op)
|
||||
op_grad_func = call.op.get_attr("FPrimalGradient")
|
||||
|
||||
# The parameter Var for gradient
|
||||
grad_var = relax.Var("grad", _numpy_to_ty(weights))
|
||||
|
||||
# Gradient mod
|
||||
grad_bb = relax.BlockBuilder()
|
||||
with grad_bb.function(func_name, param_vars + [grad_var]):
|
||||
with grad_bb.dataflow():
|
||||
orig = grad_bb.emit(call)
|
||||
# op_grad_func returns a list of Exprs representing the gradients
|
||||
grad_call = op_grad_func(orig, call, grad_var, grad_bb)
|
||||
|
||||
# Check ignore_grads
|
||||
for i, grad in enumerate(grad_call):
|
||||
if i in ignore_grads:
|
||||
assert _is_call_no_grad(grad), f"The {i}-th gradient should be no_grad"
|
||||
else:
|
||||
assert not _is_call_no_grad(grad), f"The {i}-th gradient should not be no_grad"
|
||||
|
||||
if tuple_input:
|
||||
# If the input is a tuple, the gradient is also a tuple.
|
||||
# The gradient tuple is the first (the only) element of grad_call.
|
||||
out = grad_bb.emit_output(grad_call[0])
|
||||
else:
|
||||
# We need to wrap the list into a relax.Tuple so as to emit it
|
||||
out = grad_bb.emit_output(relax.Tuple(grad_call))
|
||||
grad_bb.emit_func_output(out)
|
||||
|
||||
grad_mod = grad_bb.get()
|
||||
grad_ex = tvm.compile(grad_mod, target)
|
||||
grad_vm = relax.VirtualMachine(grad_ex, dev)
|
||||
|
||||
# tvm.runtime.Tensor inputs
|
||||
inputs_tvm = [_numpy_to_tvm(i) for i in inputs_numpy]
|
||||
weights_tvm = _numpy_to_tvm(weights)
|
||||
result_filtered = _tvm_to_numpy(grad_vm[func_name](*inputs_tvm, weights_tvm), ignore_grads)
|
||||
|
||||
# Inputs contained in ignore_grads are removed
|
||||
inputs_filtered = [inputs_numpy[i] for i in range(len(inputs_numpy)) if i not in ignore_grads]
|
||||
|
||||
check_numerical_grads(forward, inputs_filtered, result_filtered)
|
||||
|
||||
|
||||
##################### Unary #####################
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"unary_op_func,can_be_neg",
|
||||
[
|
||||
(relax.op.abs, True),
|
||||
(relax.op.cos, True),
|
||||
(relax.op.exp, True),
|
||||
(relax.op.log, False),
|
||||
(relax.op.negative, True),
|
||||
(relax.op.sigmoid, True),
|
||||
(relax.op.sin, True),
|
||||
(relax.op.sqrt, False),
|
||||
(relax.op.tanh, True),
|
||||
],
|
||||
)
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_unary(unary_op_func, can_be_neg):
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
(low, high) = (-1, 1) if can_be_neg else (0.1, 1)
|
||||
data_numpy = np.random.uniform(low, high, (3, 3)).astype(np.float32)
|
||||
relax_check_gradients(unary_op_func, [data_numpy], target, dev)
|
||||
|
||||
|
||||
##################### Binary #####################
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"binary_arith_op_func",
|
||||
[
|
||||
relax.op.add,
|
||||
relax.op.subtract,
|
||||
relax.op.multiply,
|
||||
relax.op.divide,
|
||||
relax.op.power,
|
||||
],
|
||||
)
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_binary_arith(binary_arith_op_func):
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data1_numpy = np.random.uniform(1, 2, (3, 3)).astype(np.float32)
|
||||
data2_numpy = np.random.uniform(1, 2, (3, 3)).astype(np.float32)
|
||||
relax_check_gradients(binary_arith_op_func, [data1_numpy, data2_numpy], target, dev)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("binary_minmax_op_func", [relax.op.maximum, relax.op.minimum])
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_binary_minmax(binary_minmax_op_func):
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
# Checking numerical gradient of min and max requires data1_numpy[i] != data2_numpy[i]
|
||||
# for all possible i.
|
||||
# If data1_numpy[i] == data2_numpy[i], the operator is not differentiable w.r.t. place i
|
||||
data1_numpy = np.random.uniform(1, 1.1, (3, 3)).astype(np.float32)
|
||||
delta = np.random.uniform(1, 1.1, (3, 3)).astype(np.float32)
|
||||
sign = np.random.randint(0, 2, (3, 3)).astype(np.float32) * 2 - 1
|
||||
data2_numpy = data1_numpy + delta * sign
|
||||
relax_check_gradients(binary_minmax_op_func, [data1_numpy, data2_numpy], target, dev)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"binary_cmp_op_func",
|
||||
[
|
||||
relax.op.equal,
|
||||
relax.op.greater,
|
||||
relax.op.greater_equal,
|
||||
relax.op.less,
|
||||
relax.op.less_equal,
|
||||
relax.op.not_equal,
|
||||
],
|
||||
)
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_binary_cmp(binary_cmp_op_func):
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data1_numpy = np.random.uniform(1, 2, (3, 3)).astype(np.float32)
|
||||
data2_numpy = np.random.uniform(1, 2, (3, 3)).astype(np.float32)
|
||||
relax_check_gradients(
|
||||
binary_cmp_op_func, [data1_numpy, data2_numpy], target, dev, ignore_grads=[0, 1]
|
||||
)
|
||||
|
||||
|
||||
##################### Create #####################
|
||||
|
||||
|
||||
@pytest.mark.parametrize("like_op_func", [relax.op.zeros_like, relax.op.ones_like])
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_ones_zeros_like(like_op_func):
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data_numpy = np.random.uniform(-1, 1, (3, 3)).astype(np.float32)
|
||||
relax_check_gradients(like_op_func, [data_numpy], target, dev, ignore_grads=[0])
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_full_like():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data_numpy = np.random.uniform(-1, 1, (3, 3)).astype(np.float32)
|
||||
fill_value = np.random.uniform(-1, 1, ()).astype(np.float32)
|
||||
relax_check_gradients(
|
||||
relax.op.full_like, [data_numpy, fill_value], target, dev, ignore_grads=[0, 1]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("create_op_func", [relax.op.zeros, relax.op.ones])
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_ones_zeros(create_op_func):
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
relax_check_gradients(
|
||||
create_op_func, [], target, dev, ignore_grads=[0], shape=(3, 3), dtype="float32"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_triu():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data_numpy = np.random.uniform(-1, 1, (3, 3)).astype(np.float32)
|
||||
relax_check_gradients(relax.op.triu, [data_numpy], target, dev, k=0)
|
||||
|
||||
|
||||
##################### Statistical #####################
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_sum():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data1_numpy = np.random.uniform(0, 16, (3, 3)).astype(np.float32)
|
||||
relax_check_gradients(relax.op.sum, [data1_numpy], target, dev)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_sum_with_axis():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data1_numpy = np.random.uniform(0, 16, (2, 3, 4, 5)).astype(np.float32)
|
||||
relax_check_gradients(relax.op.sum, [data1_numpy], target, dev, axis=[1, 3])
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_sum_keepdims():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data1_numpy = np.random.uniform(0, 16, (3, 3)).astype(np.float32)
|
||||
relax_check_gradients(relax.op.sum, [data1_numpy], target, dev, keepdims=True, axis=1)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_mean():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data1_numpy = np.random.uniform(0, 16, (3, 3)).astype(np.float32)
|
||||
relax_check_gradients(relax.op.mean, [data1_numpy], target, dev)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_mean_with_axis():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data1_numpy = np.random.uniform(0, 16, (2, 3, 4, 5)).astype(np.float32)
|
||||
relax_check_gradients(relax.op.mean, [data1_numpy], target, dev, axis=[1, 3])
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_mean_keepdims():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data1_numpy = np.random.uniform(0, 16, (3, 3)).astype(np.float32)
|
||||
relax_check_gradients(relax.op.mean, [data1_numpy], target, dev, keepdims=True, axis=1)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_variance():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data1_numpy = np.random.uniform(0, 16, (3, 3)).astype(np.float32)
|
||||
relax_check_gradients(relax.op.variance, [data1_numpy], target, dev)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_variance_with_axis():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data1_numpy = np.random.uniform(0, 16, (2, 3, 4, 5)).astype(np.float32)
|
||||
relax_check_gradients(relax.op.variance, [data1_numpy], target, dev, axis=[1, 3])
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_variance_keepdims():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data1_numpy = np.random.uniform(0, 16, (3, 3)).astype(np.float32)
|
||||
relax_check_gradients(relax.op.variance, [data1_numpy], target, dev, keepdims=True, axis=1)
|
||||
|
||||
|
||||
##################### Manipulate #####################
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_reshape():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data_numpy = np.random.uniform(0, 16, (2, 3, 5)).astype(np.float32)
|
||||
relax_check_gradients(
|
||||
relax.op.reshape, [data_numpy], target, dev, ignore_grads=[1], shape=(5, 6)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_reshape_infer_dim():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data_numpy = np.random.uniform(0, 16, (2, 3, 5)).astype(np.float32)
|
||||
relax_check_gradients(
|
||||
relax.op.reshape, [data_numpy], target, dev, ignore_grads=[1], shape=(5, 2, 1, -1)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_permute_dims():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data_numpy = np.random.uniform(0, 16, (2, 3, 4, 5)).astype(np.float32)
|
||||
relax_check_gradients(relax.op.permute_dims, [data_numpy], target, dev)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_permute_dims_with_axes():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data_numpy = np.random.uniform(0, 16, (2, 3, 4, 5)).astype(np.float32)
|
||||
relax_check_gradients(
|
||||
relax.op.permute_dims,
|
||||
[data_numpy],
|
||||
target,
|
||||
dev,
|
||||
axes=(0, 3, 1, 2),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_concat():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data_numpy1 = np.random.uniform(1, 16, (3, 3)).astype(np.float32)
|
||||
data_numpy2 = np.random.uniform(1, 16, (3, 4)).astype(np.float32)
|
||||
data_numpy3 = np.random.uniform(1, 16, (3, 5)).astype(np.float32)
|
||||
relax_check_gradients(
|
||||
relax.op.concat,
|
||||
[data_numpy1, data_numpy2, data_numpy3],
|
||||
target,
|
||||
dev,
|
||||
tuple_input=True,
|
||||
axis=1,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_split_indices():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data_numpy = np.random.uniform(1, 16, (3, 12)).astype(np.float32)
|
||||
relax_check_gradients(
|
||||
relax.op.split,
|
||||
[data_numpy],
|
||||
target,
|
||||
dev,
|
||||
indices_or_sections=[3, 7],
|
||||
axis=1,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_split_section():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data_numpy = np.random.uniform(1, 16, (3, 12)).astype(np.float32)
|
||||
relax_check_gradients(
|
||||
relax.op.split,
|
||||
[data_numpy],
|
||||
target,
|
||||
dev,
|
||||
indices_or_sections=3,
|
||||
axis=1,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_reshape():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data_numpy = np.random.uniform(1, 16, (3, 4)).astype(np.float32)
|
||||
|
||||
relax_check_gradients(
|
||||
relax.op.reshape,
|
||||
[data_numpy],
|
||||
target,
|
||||
dev,
|
||||
shape=(3, 2, 2),
|
||||
ignore_grads=[1],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_cumsum():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data_numpy1 = np.random.uniform(1, 16, (3, 3)).astype(np.float32)
|
||||
relax_check_gradients(
|
||||
relax.op.cumsum,
|
||||
[data_numpy1],
|
||||
target,
|
||||
dev,
|
||||
axis=1,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_cumsum_no_axis():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data_numpy1 = np.random.uniform(1, 16, (3, 3)).astype(np.float32)
|
||||
relax_check_gradients(
|
||||
relax.op.cumsum,
|
||||
[data_numpy1],
|
||||
target,
|
||||
dev,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_expand_dims():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data_numpy = np.random.uniform(1, 16, (3, 12)).astype(np.float32)
|
||||
relax_check_gradients(relax.op.expand_dims, [data_numpy], target, dev, axis=1)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_expand_dims_list():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data_numpy = np.random.uniform(1, 16, (3, 12)).astype(np.float32)
|
||||
relax_check_gradients(relax.op.expand_dims, [data_numpy], target, dev, axis=(0, 2, 3))
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_broadcast_to():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data_numpy = np.random.uniform(1, 16, (3, 4)).astype(np.float32)
|
||||
relax_check_gradients(
|
||||
relax.op.broadcast_to,
|
||||
[data_numpy],
|
||||
target,
|
||||
dev,
|
||||
shape=(2, 3, 4),
|
||||
ignore_grads=[1],
|
||||
)
|
||||
|
||||
|
||||
##################### Index #####################
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_take():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data_numpy = np.random.uniform(0, 16, size=(2, 3, 4)).astype(np.float32)
|
||||
indices = np.array([0, 1])
|
||||
relax_check_gradients(
|
||||
relax.op.take,
|
||||
[data_numpy, indices],
|
||||
target,
|
||||
dev,
|
||||
axis=1,
|
||||
ignore_grads=[1],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_take_no_axis():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data_numpy = np.random.uniform(0, 16, size=(5,)).astype(np.float32)
|
||||
indices = np.array([1, 3])
|
||||
relax_check_gradients(
|
||||
relax.op.take,
|
||||
[data_numpy, indices],
|
||||
target,
|
||||
dev,
|
||||
ignore_grads=[1],
|
||||
)
|
||||
|
||||
|
||||
##################### Search #####################
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_where():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data1_numpy = np.random.uniform(0, 1, size=(3, 3)) > 0.5
|
||||
data2_numpy = np.random.uniform(0, 16, size=(3, 3)).astype(np.float32)
|
||||
data3_numpy = np.random.uniform(0, 16, size=(3, 3)).astype(np.float32)
|
||||
|
||||
relax_check_gradients(
|
||||
relax.op.where,
|
||||
[data1_numpy, data2_numpy, data3_numpy],
|
||||
target,
|
||||
dev,
|
||||
ignore_grads=[0],
|
||||
)
|
||||
|
||||
|
||||
##################### Linear Algebra #####################
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_matmul_2_2():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data1_numpy = np.random.uniform(0, 16, (2, 3)).astype(np.float32)
|
||||
data2_numpy = np.random.uniform(0, 16, (3, 4)).astype(np.float32)
|
||||
relax_check_gradients(relax.op.matmul, [data1_numpy, data2_numpy], target, dev)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_matmul_1_1():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data1_numpy = np.random.uniform(0, 16, (4,)).astype(np.float32)
|
||||
data2_numpy = np.random.uniform(0, 16, (4,)).astype(np.float32)
|
||||
relax_check_gradients(relax.op.matmul, [data1_numpy, data2_numpy], target, dev)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_matmul_1_4():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data1_numpy = np.random.uniform(0, 16, (4,)).astype(np.float32)
|
||||
data2_numpy = np.random.uniform(0, 16, (2, 3, 4, 5)).astype(np.float32)
|
||||
relax_check_gradients(relax.op.matmul, [data1_numpy, data2_numpy], target, dev)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_matmul_4_1():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data1_numpy = np.random.uniform(0, 16, (2, 3, 4, 5)).astype(np.float32)
|
||||
data2_numpy = np.random.uniform(0, 16, (5,)).astype(np.float32)
|
||||
relax_check_gradients(relax.op.matmul, [data1_numpy, data2_numpy], target, dev)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_matmul_5_4():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data1_numpy = np.random.uniform(0, 16, (2, 3, 1, 4, 5)).astype(np.float32)
|
||||
data2_numpy = np.random.uniform(0, 16, (3, 2, 5, 4)).astype(np.float32)
|
||||
relax_check_gradients(
|
||||
relax.op.matmul,
|
||||
[data1_numpy, data2_numpy],
|
||||
target,
|
||||
dev,
|
||||
)
|
||||
|
||||
|
||||
##################### Datatype #####################
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_astype():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data_numpy = np.random.uniform(0, 16, size=(3, 3)).astype(np.float64)
|
||||
relax_check_gradients(relax.op.astype, [data_numpy], target, dev, dtype="float32")
|
||||
|
||||
|
||||
##################### Neural network #####################
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_relu():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data1_numpy = np.random.uniform(0.2, 1, (3, 3)).astype(np.float32)
|
||||
sign = np.random.randint(0, 2, (3, 3)).astype(np.float32) * 2 - 1
|
||||
data1_numpy *= sign
|
||||
relax_check_gradients(relax.op.nn.relu, [data1_numpy], target, dev)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_silu():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data1_numpy = np.random.uniform(0, 16, (3, 3)).astype(np.float32)
|
||||
relax_check_gradients(relax.op.nn.silu, [data1_numpy], target, dev)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_softmax():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data1_numpy = np.random.uniform(0, 16, (3, 3)).astype(np.float32)
|
||||
relax_check_gradients(relax.op.nn.softmax, [data1_numpy], target, dev)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_softmax_with_axis():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data1_numpy = np.random.uniform(0, 16, (3, 3)).astype(np.float32)
|
||||
relax_check_gradients(relax.op.nn.softmax, [data1_numpy], target, dev, axis=1)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_log_softmax():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data1_numpy = np.random.uniform(0, 16, (3, 3)).astype(np.float32)
|
||||
relax_check_gradients(relax.op.nn.log_softmax, [data1_numpy], target, dev)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_log_softmax_with_axis():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data1_numpy = np.random.uniform(0, 16, (3, 3)).astype(np.float32)
|
||||
relax_check_gradients(relax.op.nn.log_softmax, [data1_numpy], target, dev, axis=1)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_cross_entropy_with_logits():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data_numpy1 = np.random.uniform(1, 16, (3,)).astype(np.float32)
|
||||
data_numpy2 = np.random.uniform(1, 16, (3,)).astype(np.float32)
|
||||
relax_check_gradients(
|
||||
relax.op.nn.cross_entropy_with_logits,
|
||||
[data_numpy1, data_numpy2],
|
||||
target,
|
||||
dev,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_cross_entropy_with_logits_batch():
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data_numpy1 = np.random.uniform(1, 16, (2, 3)).astype(np.float32)
|
||||
data_numpy2 = np.random.uniform(1, 16, (2, 3)).astype(np.float32)
|
||||
relax_check_gradients(
|
||||
relax.op.nn.cross_entropy_with_logits,
|
||||
[data_numpy1, data_numpy2],
|
||||
target,
|
||||
dev,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"nll_reduction,nll_weighted,nll_ignore_index",
|
||||
[
|
||||
("mean", True, -1),
|
||||
("sum", True, -1),
|
||||
("none", True, -1),
|
||||
("mean", True, 1),
|
||||
("mean", True, 1),
|
||||
("mean", False, 1),
|
||||
],
|
||||
)
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_nll_loss(nll_reduction, nll_weighted, nll_ignore_index):
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data1_numpy = np.random.uniform(0, 16, (2, 3, 4)).astype(np.float32)
|
||||
data2_numpy = np.random.randint(0, 3, (2, 4)).astype(np.int64)
|
||||
# force a position in targets it not ignore_index, to avoid zero total weight
|
||||
data2_numpy[0][0] = 0
|
||||
# weight > 0
|
||||
data3_numpy = np.random.uniform(1, 16, (3,)).astype(np.float32)
|
||||
|
||||
input = [data1_numpy, data2_numpy] + ([data3_numpy] if nll_weighted else [])
|
||||
ignore_grads = [1] + ([2] if nll_weighted else [])
|
||||
|
||||
relax_check_gradients(
|
||||
relax.op.nn.nll_loss,
|
||||
input,
|
||||
target,
|
||||
dev,
|
||||
ignore_grads=ignore_grads,
|
||||
reduction=nll_reduction,
|
||||
ignore_index=nll_ignore_index,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"nll_reduction1,nll_weighted1,nll_ignore_index1",
|
||||
[
|
||||
("mean", True, -1),
|
||||
("sum", True, -1),
|
||||
("none", True, -1),
|
||||
],
|
||||
)
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_nll_loss_no_batch(nll_reduction1, nll_weighted1, nll_ignore_index1):
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data1_numpy = np.random.uniform(0, 16, (3,)).astype(np.float32)
|
||||
data2_numpy = np.random.randint(0, 3, ()).astype(np.int64)
|
||||
# weight > 0
|
||||
data3_numpy = np.random.uniform(1, 16, (3,)).astype(np.float32)
|
||||
|
||||
input = [data1_numpy, data2_numpy] + ([data3_numpy] if nll_weighted1 else [])
|
||||
ignore_grads = [1] + ([2] if nll_weighted1 else [])
|
||||
|
||||
relax_check_gradients(
|
||||
relax.op.nn.nll_loss,
|
||||
input,
|
||||
target,
|
||||
dev,
|
||||
ignore_grads=ignore_grads,
|
||||
reduction=nll_reduction1,
|
||||
ignore_index=nll_ignore_index1,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"c2d_shape1,c2d_shape2,c2d_kwargs",
|
||||
[
|
||||
(
|
||||
(3, 2, 10, 10),
|
||||
(3, 2, 3, 3),
|
||||
{},
|
||||
),
|
||||
(
|
||||
(3, 2, 10, 10),
|
||||
(3, 2, 1, 2),
|
||||
{},
|
||||
),
|
||||
(
|
||||
(3, 2, 10, 10),
|
||||
(3, 2, 3, 3),
|
||||
{"strides": (2, 2), "padding": (3, 2), "dilation": (1, 1)},
|
||||
),
|
||||
(
|
||||
(3, 2, 10, 10),
|
||||
(3, 2, 3, 3),
|
||||
{"strides": (2, 1), "padding": (2, 2), "dilation": (1, 1)},
|
||||
),
|
||||
(
|
||||
(3, 6, 10, 10),
|
||||
(4, 3, 3, 3),
|
||||
{"groups": 2},
|
||||
),
|
||||
(
|
||||
(3, 2, 10, 10),
|
||||
(4, 1, 3, 3),
|
||||
{"groups": 2, "strides": (2, 2), "padding": (2, 2), "dilation": (1, 1)},
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_conv2d(c2d_shape1, c2d_shape2, c2d_kwargs):
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
import pytest
|
||||
|
||||
# Use smaller range to reduce numerical errors in gradient check
|
||||
data1_numpy = np.random.uniform(0, 2, c2d_shape1).astype(np.float32)
|
||||
data2_numpy = np.random.uniform(0, 2, c2d_shape2).astype(np.float32)
|
||||
relax_check_gradients(
|
||||
relax.op.nn.conv2d,
|
||||
[data1_numpy, data2_numpy],
|
||||
target,
|
||||
dev,
|
||||
**c2d_kwargs,
|
||||
)
|
||||
|
||||
|
||||
pool_params = [
|
||||
(
|
||||
(3, 3),
|
||||
{},
|
||||
),
|
||||
(
|
||||
(3, 3),
|
||||
{"strides": (2, 2), "padding": (1, 2), "dilation": (1, 1), "count_include_pad": True},
|
||||
),
|
||||
(
|
||||
(5, 5),
|
||||
{
|
||||
"strides": (2, 2),
|
||||
"padding": (2, 1),
|
||||
"dilation": (1, 1),
|
||||
"ceil_mode": True,
|
||||
"count_include_pad": True,
|
||||
},
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pool_size,pool_kwargs", pool_params)
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_max_pool2d(pool_size, pool_kwargs):
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data_numpy = np.random.uniform(0, 3, size=(3, 2, 10, 10)).astype(np.float32)
|
||||
relax_check_gradients(
|
||||
relax.op.nn.max_pool2d,
|
||||
[data_numpy],
|
||||
target,
|
||||
dev,
|
||||
pool_size=pool_size,
|
||||
**pool_kwargs,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pool_size,pool_kwargs", pool_params)
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
||||
def test_avg_pool2d(pool_size, pool_kwargs):
|
||||
target = "llvm"
|
||||
dev = tvm.device(target)
|
||||
data_numpy = np.random.uniform(0, 3, size=(3, 2, 10, 10)).astype(np.float32)
|
||||
relax_check_gradients(
|
||||
relax.op.nn.avg_pool2d,
|
||||
[data_numpy],
|
||||
target,
|
||||
dev,
|
||||
pool_size=pool_size,
|
||||
**pool_kwargs,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
Reference in New Issue
Block a user