()
| 266 | |
| 267 | |
| 268 | def test_elemwise_relu(): |
| 269 | x_np = [1.0, -1.0] |
| 270 | dz_np = [1.0] |
| 271 | x = mge.Tensor(x_np) |
| 272 | dz = mge.Tensor(dz_np) |
| 273 | |
| 274 | refs = {} |
| 275 | |
| 276 | def f(x): |
| 277 | x = x * 2 |
| 278 | refs["x"] = TensorWeakRef(x) |
| 279 | return relu(x) |
| 280 | |
| 281 | with Grad() as grad: |
| 282 | grad.wrt(x, callback=save_to(x)) |
| 283 | z = f(x) |
| 284 | assert refs["x"]() is None |
| 285 | grad(z, dz) |
| 286 | |
| 287 | np.testing.assert_almost_equal(x.grad.numpy(), [2.0, 0]) |
| 288 | |
| 289 | |
| 290 | def test_elemwise_relu_backward_fn(): |