diff --git a/CMakeLists.txt b/CMakeLists.txt index 363b2056a87a..8765a3346069 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -83,7 +83,7 @@ else(MSVC) include(CheckCXXCompilerFlag) check_cxx_compiler_flag("-std=c++11" SUPPORT_CXX11) if ("${CMAKE_BUILD_TYPE}" STREQUAL "Debug") - add_compile_options(-Wall -fPIC -std=c++11) + add_compile_options(-O0 -Wall -fPIC -std=c++11) else() set(CMAKE_C_FLAGS "-O2 -Wall -fPIC ${CMAKE_C_FLAGS}") set(CMAKE_CXX_FLAGS "-O2 -Wall -fPIC -std=c++11 ${CMAKE_CXX_FLAGS}") diff --git a/python/tvm/relay/op/_tensor.py b/python/tvm/relay/op/_tensor.py index 240e7fffd328..d9b5e2e89ce0 100644 --- a/python/tvm/relay/op/_tensor.py +++ b/python/tvm/relay/op/_tensor.py @@ -5,17 +5,21 @@ from .op import register_compute, register_schedule, register_pattern from .op import register_gradient from .op import schedule_injective, OpPattern +from .transform import collapse_sum_like +from .tensor import negative + def add_grad(orig, grad): - from tvm.relay import op - return [op.broadcast_to_like(grad, orig.args[0]), op.broadcast_to_like(grad, orig.args[1])] + return [collapse_sum_like(grad, orig.args[0]), collapse_sum_like(grad, orig.args[1])] + register_gradient("add", add_grad) + def subtract_grad(orig, grad): - from tvm.relay import op - return [op.broadcast_to_like(grad, orig.args[0]), - op.broadcast_to_like(op.negative(grad), orig.args[1])] + return [collapse_sum_like(grad, orig.args[0]), + collapse_sum_like(negative(grad), orig.args[1])] + register_gradient("subtract", subtract_grad) diff --git a/tests/python/relay/test_ad.py b/tests/python/relay/test_ad.py index 7844236907c4..6b5d0e776934 100644 --- a/tests/python/relay/test_ad.py +++ b/tests/python/relay/test_ad.py @@ -69,8 +69,64 @@ def test_sub(): np.testing.assert_allclose(grad.asnumpy(), np.zeros_like(x.asnumpy())) +def test_broadcast_add(): + shape1 = (3, 4, 1) + shape2 = (1, 5) + dtype = 'float32' + x_nd = rand(dtype, *shape1) + y_nd = rand(dtype, *shape2) + x_np = x_nd.asnumpy() + y_np = y_nd.asnumpy() + expected_forward = x_np + y_np + t1 = relay.TensorType(shape1, dtype) + t2 = relay.TensorType(shape2, dtype) + x = relay.var("x", t1) + y = relay.var("y", t2) + func = relay.Function([x, y], x + y) + full_func = relay.ir_pass.infer_type(gradient(func)) + assert full_func.checked_type == relay.FuncType([t1, t2], + relay.TupleType([relay.TensorType(expected_forward.shape, dtype), + relay.TupleType([t1, t2])])) + ex = create_executor() + forward, (grad_x, grad_y) = ex.evaluate(full_func)(x_nd, y_nd) + np.testing.assert_allclose(forward.asnumpy(), expected_forward) + np.testing.assert_allclose(grad_x.asnumpy(), + np.ones_like(expected_forward).sum(axis=2, keepdims=True)) + np.testing.assert_allclose(grad_y.asnumpy(), + np.ones_like(expected_forward).sum(axis=(0, 1), keepdims=True).squeeze(axis=0)) + + +def test_broadcast_subtract(): + shape1 = (3, 4, 1) + shape2 = (1, 5) + dtype = 'float32' + x_nd = rand(dtype, *shape1) + y_nd = rand(dtype, *shape2) + x_np = x_nd.asnumpy() + y_np = y_nd.asnumpy() + expected_forward = x_np - y_np + t1 = relay.TensorType(shape1, dtype) + t2 = relay.TensorType(shape2, dtype) + x = relay.var("x", t1) + y = relay.var("y", t2) + func = relay.Function([x, y], x - y) + full_func = relay.ir_pass.infer_type(gradient(func)) + assert full_func.checked_type == relay.FuncType([t1, t2], + relay.TupleType([relay.TensorType(expected_forward.shape, dtype), + relay.TupleType([t1, t2])])) + ex = create_executor() + forward, (grad_x, grad_y) = ex.evaluate(full_func)(x_nd, y_nd) + np.testing.assert_allclose(forward.asnumpy(), expected_forward) + np.testing.assert_allclose(grad_x.asnumpy(), + np.ones_like(expected_forward).sum(axis=2, keepdims=True)) + np.testing.assert_allclose(grad_y.asnumpy(), + -np.ones_like(expected_forward).sum(axis=(0, 1), keepdims=True).squeeze(axis=0)) + + if __name__ == "__main__": test_id() test_add() test_temp_add() test_sub() + test_broadcast_add() + test_broadcast_subtract()