Function test_elemwise_add

imperative/python/test/unit/core/test_autodiff.py:241–265

()

def test_elemwise_add():
    x_np = np.random.rand(10).astype("float32")
    y_np = np.random.rand(10, 10).astype("float32")
    dz_np = np.random.rand(10, 10).astype("float32")
    x = mge.Tensor(x_np)
    y = mge.Tensor(y_np)
    dz = mge.Tensor(dz_np)

    refs = {}

    def f(x, y):
        x = x * 2
        refs["x"] = TensorWeakRef(x)
        refs["y"] = TensorWeakRef(y)
        return x + y

    with Grad() as grad:
        grad.wrt(x, callback=save_to(x))
        z = f(x, y)
        del y
        for k, r in refs.items():
            assert r() is None
        grad(z, dz)

    np.testing.assert_almost_equal(x.grad.numpy(), dz_np.sum(0) * 2, decimal=5)


def test_elemwise_relu():