()
| 480 | |
| 481 | |
| 482 | def test_addAxis(): |
| 483 | x_np = np.random.rand(3, 3).astype("float32") |
| 484 | x = mge.Tensor(x_np) |
| 485 | |
| 486 | with Grad() as grad: |
| 487 | grad.wrt(x, callback=save_to(x)) |
| 488 | refs = {} |
| 489 | |
| 490 | def f(x): |
| 491 | x = x * 1 |
| 492 | y = F.expand_dims(x, [2, 3]) |
| 493 | refs["x"] = TensorWeakRef(x) |
| 494 | return y |
| 495 | |
| 496 | y = f(x) |
| 497 | for _, r in refs.items(): |
| 498 | assert r() is None |
| 499 | grad(y, F.ones_like(y)) |
| 500 | |
| 501 | np.testing.assert_equal(np.ones((3, 3), dtype=np.float32), x.grad.numpy()) |
| 502 | |
| 503 | |
| 504 | def test_removeAxis(): |