Function test_AxisAddRemove

imperative/python/test/unit/core/test_autodiff.py:410–431

()

def test_AxisAddRemove():
    x_np = np.random.rand(1, 5).astype("float32")
    x = mge.Tensor(x_np)

    with Grad() as grad:
        grad.wrt(x, callback=save_to(x))
        refs = {}

        def f(x):
            x = x * 1
            y = F.squeeze(F.expand_dims(x, 2), 0)
            refs["x"] = TensorWeakRef(x)
            return y

        y = f(x)
        for _, r in refs.items():
            assert r() is None
        grad(y, F.ones_like(y))

    np.testing.assert_equal(
        np.array([[1, 1, 1, 1, 1]], dtype=np.float32), x.grad.numpy()
    )


def test_Broadcast():