def test_relu_gradient(): shape = (2000, 2500) x = np.random.uniform(-1, 1, shape).astype(dtype) grad_x = np.random.uniform(-5, 5, shape).astype(dtype) y = np.zeros(shape).astype(dtype) arr_x = tvm.nd.array(x, ctx=ctx) arr_grad_x = tvm.nd.array(grad_x, ctx=ctx) arr_y = tvm.nd.array(y, ctx=ctx) relu_gradient = tvm_op.make_relu_gradient(shape, tgt, tgt_host, "relu_gradient") relu_gradient(arr_x, arr_grad_x, arr_y) y = arr_y.asnumpy() np.testing.assert_allclose(((x > 0) * grad_x).astype(dtype), y)