(self)
| 628 | self.assertTrue(str(inorm) == "InstanceNorm(3, eps=1e-05, affine=False)") |
| 629 | |
| 630 | def test_batch_norm(self): |
| 631 | mx.random.seed(42) |
| 632 | x = mx.random.normal((5, 4), dtype=mx.float32) |
| 633 | |
| 634 | # Batch norm |
| 635 | bn = nn.BatchNorm(num_features=4, affine=True) |
| 636 | self.assertTrue(mx.allclose(bn.running_mean, mx.zeros_like(bn.running_mean))) |
| 637 | self.assertTrue(mx.allclose(bn.running_var, mx.ones_like(bn.running_var))) |
| 638 | y = bn(x) |
| 639 | expected_y = mx.array( |
| 640 | [ |
| 641 | [-0.439520, 1.647328, -0.955515, 1.966031], |
| 642 | [-1.726690, -1.449826, -0.234026, -0.723364], |
| 643 | [0.938414, -0.349603, -0.354470, -0.175369], |
| 644 | [0.305006, 0.234914, -0.393017, -0.459385], |
| 645 | [0.922789, -0.082813, 1.937028, -0.607913], |
| 646 | ], |
| 647 | ) |
| 648 | expected_mean = mx.array([0.008929, 0.005680, -0.016092, 0.027778]) |
| 649 | expected_var = mx.array([0.928435, 1.00455, 1.04117, 0.94258]) |
| 650 | self.assertTrue(x.shape == y.shape) |
| 651 | self.assertTrue(mx.allclose(y, expected_y, atol=1e-5)) |
| 652 | self.assertTrue(mx.allclose(bn.running_mean, expected_mean, atol=1e-5)) |
| 653 | self.assertTrue(mx.allclose(bn.running_var, expected_var, atol=1e-5)) |
| 654 | |
| 655 | # test eval mode |
| 656 | bn.eval() |
| 657 | y = bn(x) |
| 658 | expected_y = mx.array( |
| 659 | [ |
| 660 | [-0.15984, 1.73159, -1.25456, 1.57891], |
| 661 | [-0.872193, -1.4281, -0.414439, -0.228678], |
| 662 | [0.602743, -0.30566, -0.554687, 0.139639], |
| 663 | [0.252199, 0.29066, -0.599572, -0.0512532], |
| 664 | [0.594096, -0.0334829, 2.11359, -0.151081], |
| 665 | ] |
| 666 | ) |
| 667 | |
| 668 | self.assertTrue(x.shape == y.shape) |
| 669 | self.assertTrue(mx.allclose(y, expected_y, atol=1e-5)) |
| 670 | |
| 671 | # test_no_affine |
| 672 | bn = nn.BatchNorm(num_features=4, affine=False) |
| 673 | y = bn(x) |
| 674 | expected_y = mx.array( |
| 675 | [ |
| 676 | [-0.439520, 1.647328, -0.955515, 1.966031], |
| 677 | [-1.726690, -1.449826, -0.234026, -0.723364], |
| 678 | [0.938414, -0.349603, -0.354470, -0.175369], |
| 679 | [0.305006, 0.234914, -0.393017, -0.459385], |
| 680 | [0.922789, -0.082813, 1.937028, -0.607913], |
| 681 | ] |
| 682 | ) |
| 683 | self.assertTrue(x.shape == y.shape) |
| 684 | self.assertTrue(mx.allclose(y, expected_y, atol=1e-5)) |
| 685 | |
| 686 | # test with 3D input |
| 687 | mx.random.seed(42) |
nothing calls this directly
no test coverage detected