| 410 | self.assertTrue(np.allclose(var, 4 * np.ones_like(var), atol=1e-6)) |
| 411 | |
| 412 | def test_instance_norm(self): |
| 413 | # Test InstanceNorm1d |
| 414 | x = mx.array( |
| 415 | [ |
| 416 | [ |
| 417 | [-0.0119524, 1.1263, 2.02223], |
| 418 | [-0.500331, 0.517899, -1.21143], |
| 419 | [1.12958, -0.21413, -2.48738], |
| 420 | [1.39955, 0.891329, 1.63289], |
| 421 | ], |
| 422 | [ |
| 423 | [0.241417, -0.619157, -0.77484], |
| 424 | [-1.42512, 0.970817, -1.31352], |
| 425 | [2.739, -1.2506, 1.56844], |
| 426 | [-1.23175, 0.32756, 1.13969], |
| 427 | ], |
| 428 | ] |
| 429 | ) |
| 430 | inorm = nn.InstanceNorm(dims=3) |
| 431 | y = inorm(x) |
| 432 | expected_y = [ |
| 433 | [ |
| 434 | [-0.657082, 1.07593, 1.0712], |
| 435 | [-1.27879, -0.123074, -0.632505], |
| 436 | [0.796101, -1.56572, -1.30476], |
| 437 | [1.13978, 0.612862, 0.866067], |
| 438 | ], |
| 439 | [ |
| 440 | [0.0964426, -0.557906, -0.759885], |
| 441 | [-0.904772, 1.30444, -1.20013], |
| 442 | [1.59693, -1.29752, 1.15521], |
| 443 | [-0.7886, 0.550987, 0.804807], |
| 444 | ], |
| 445 | ] |
| 446 | self.assertTrue(x.shape == y.shape) |
| 447 | self.assertTrue(np.allclose(y, expected_y, atol=1e-5)) |
| 448 | # Test InstanceNorm2d |
| 449 | x = mx.array( |
| 450 | [ |
| 451 | [ |
| 452 | [ |
| 453 | [-0.458824, 0.483254, -0.58611], |
| 454 | [-0.447996, -0.176577, -0.622545], |
| 455 | [0.0486988, -0.0611224, 1.8845], |
| 456 | ], |
| 457 | [ |
| 458 | [1.13049, 0.345315, -0.926389], |
| 459 | [0.301795, 0.99207, -0.184927], |
| 460 | [-2.23876, -0.758631, -1.12639], |
| 461 | ], |
| 462 | [ |
| 463 | [0.0986325, -1.82973, -0.241765], |
| 464 | [-1.25257, 0.154442, -0.556204], |
| 465 | [-0.329399, -0.319107, 0.830584], |
| 466 | ], |
| 467 | ], |
| 468 | [ |
| 469 | [ |