| 320 | |
| 321 | |
| 322 | def test_subtensor(): |
| 323 | x_np = np.random.rand(3, 3).astype("float32") |
| 324 | x = mge.Tensor(x_np) |
| 325 | |
| 326 | with Grad() as grad: |
| 327 | grad.wrt(x, callback=save_to(x)) |
| 328 | refs = {} |
| 329 | |
| 330 | def f(x): |
| 331 | x = x * 1 |
| 332 | y = x[1:-1, :2] |
| 333 | refs["x"] = TensorWeakRef(x) |
| 334 | return y |
| 335 | |
| 336 | y = f(x) |
| 337 | for _, r in refs.items(): |
| 338 | assert r() is None |
| 339 | grad(y, F.ones_like(y)) |
| 340 | |
| 341 | np.testing.assert_equal( |
| 342 | np.array([[0, 0, 0], [1, 1, 0], [0, 0, 0]], dtype=np.float32), x.grad.numpy() |
| 343 | ) |
| 344 | |
| 345 | x_np = np.random.rand(3, 2).astype("float32") |
| 346 | x = mge.Tensor(x_np) |
| 347 | |
| 348 | with Grad() as grad: |
| 349 | grad.wrt(x, callback=save_to(x)) |
| 350 | refs = {} |
| 351 | |
| 352 | def f(x): |
| 353 | x = x * 1 |
| 354 | y = x[True, 0:1] |
| 355 | refs["x"] = TensorWeakRef(x) |
| 356 | return y |
| 357 | |
| 358 | y = f(x) |
| 359 | for _, r in refs.items(): |
| 360 | assert r() is None |
| 361 | grad(y, F.ones_like(y)) |
| 362 | |
| 363 | np.testing.assert_equal( |
| 364 | np.array([[1, 1], [0, 0], [0, 0]], dtype=np.float32), x.grad.numpy() |
| 365 | ) |
| 366 | |
| 367 | with Grad() as grad: |
| 368 | grad.wrt(x, callback=save_to(x)) |
| 369 | refs = {} |
| 370 | |
| 371 | def f(x): |
| 372 | x = x * 1 |
| 373 | y = x[False, 0:1] |
| 374 | refs["x"] = TensorWeakRef(x) |
| 375 | return y |
| 376 | |
| 377 | y = f(x) |
| 378 | for _, r in refs.items(): |
| 379 | assert r() is None |