This commit is contained in:
Awni Hannun 2023-12-11 09:26:49 -08:00 committed by GitHub
parent f4ddd7dc44
commit 69505b4e9b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 4 additions and 4 deletions

View File

@ -203,4 +203,4 @@ def _reduce(loss: mx.array, reduction: str = "none"):
elif reduction == "none":
return loss
else:
raise ValueError("Invalid reduction. Must be 'none', 'mean', or 'sum'.")
raise ValueError("Invalid reduction. Must be 'none', 'mean', or 'sum'.")

View File

@ -107,17 +107,17 @@ class TestNN(mlx_tests.MLXTestCase):
# Test with reduction 'none'
losses_none = nn.losses.binary_cross_entropy(inputs, targets, reduction="none")
expected_none = mx.array([[0.693147, 0.693147], [0.693147, 0.693147]])
self.assertTrue(mx.array_equal(losses_none, expected_none))
self.assertTrue(mx.allclose(losses_none, expected_none))
# Test with reduction 'mean'
losses_mean = nn.losses.binary_cross_entropy(inputs, targets, reduction="mean")
expected_mean = mx.mean(expected_none)
self.assertEqual(losses_mean, expected_mean)
self.assertTrue(mx.allclose(losses_mean, expected_mean))
# Test with reduction 'sum'
losses_sum = nn.losses.binary_cross_entropy(inputs, targets, reduction="sum")
expected_sum = mx.sum(expected_none)
self.assertEqual(losses_sum, expected_sum)
self.assertTrue(mx.allclose(losses_sum, expected_sum))
def test_gelu(self):
inputs = [1.15286231, -0.81037411, 0.35816911, 0.77484438, 0.66276414]