mirror of
https://github.com/ml-explore/mlx.git
synced 2025-07-16 05:41:14 +08:00
fixes (#131)
This commit is contained in:
parent
f4ddd7dc44
commit
69505b4e9b
@ -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'.")
|
||||
|
@ -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]
|
||||
|
Loading…
Reference in New Issue
Block a user