diff --git a/mlx/backend/metal/kernels/binary.metal b/mlx/backend/metal/kernels/binary.metal index e31ffcbc8..b7bbe6f37 100644 --- a/mlx/backend/metal/kernels/binary.metal +++ b/mlx/backend/metal/kernels/binary.metal @@ -18,7 +18,7 @@ struct Remainder { template T operator()(T x, T y) { return x % y; } template <> float operator()(float x, float y) { return fmod(x, y); } template <> half operator()(half x, half y) { return fmod(x, y); } - template <> bfloat operator()(bfloat x, bfloat y) { return fmod(x, y); } + template <> bfloat16_t operator()(bfloat16_t x, bfloat16_t y) { return fmod(x, y); } }; struct Equal { diff --git a/mlx/backend/metal/kernels/complex.h b/mlx/backend/metal/kernels/complex.h index 097a7e3f8..6bd427bf8 100644 --- a/mlx/backend/metal/kernels/complex.h +++ b/mlx/backend/metal/kernels/complex.h @@ -112,5 +112,7 @@ constexpr complex64_t operator*(complex64_t a, complex64_t b) { } constexpr complex64_t operator%(complex64_t a, complex64_t b) { - return {fmod(a.real, b.real), fmod(a.imag, b.imag)}; + auto real = a.real - (b.real * static_cast(a.real / b.real)); + auto imag = a.imag - (b.imag * static_cast(a.imag / b.imag)); + return {real, imag}; } diff --git a/python/mlx/nn/losses.py b/python/mlx/nn/losses.py index 068d2db78..4e0d14ef7 100644 --- a/python/mlx/nn/losses.py +++ b/python/mlx/nn/losses.py @@ -2,6 +2,7 @@ import mlx.core as mx + def cross_entropy( logits: mx.array, targets: mx.array, axis: int = -1, reduction: str = "none" ) -> mx.array: @@ -24,7 +25,9 @@ def cross_entropy( return _reduce(loss, reduction) -def l1_loss(predictions: mx.array, targets: mx.array, reduction: str = "none") -> mx.array: +def l1_loss( + predictions: mx.array, targets: mx.array, reduction: str = "none" +) -> mx.array: """ Computes the L1 loss between predictions and targets. @@ -38,11 +41,13 @@ def l1_loss(predictions: mx.array, targets: mx.array, reduction: str = "none") - mx.array: The computed L1 loss. """ loss = mx.mean(mx.abs(predictions - targets)) - + return _reduce(loss, reduction) -def mse_loss(predictions: mx.array, targets: mx.array, axis: int = -1, reduction: str = "none") -> mx.array: +def mse_loss( + predictions: mx.array, targets: mx.array, axis: int = -1, reduction: str = "none" +) -> mx.array: """ Computes the mean squared error loss between predictions and targets. @@ -57,11 +62,13 @@ def mse_loss(predictions: mx.array, targets: mx.array, axis: int = -1, reduction mx.array: The computed mean squared error loss. """ loss = mx.mean(mx.square(predictions - targets), axis) - + return _reduce(loss, reduction) -def nll_loss(logits: mx.array, targets: mx.array, axis: int = -1, reduction: str = "none") -> mx.array: +def nll_loss( + logits: mx.array, targets: mx.array, axis: int = -1, reduction: str = "none" +) -> mx.array: """ Computes the negative log likelihood loss between logits and targets. @@ -80,7 +87,9 @@ def nll_loss(logits: mx.array, targets: mx.array, axis: int = -1, reduction: str return _reduce(loss, reduction) -def kl_div_loss(logits: mx.array, targets: mx.array, axis: int = -1, reduction: str = "none") -> mx.array: +def kl_div_loss( + logits: mx.array, targets: mx.array, axis: int = -1, reduction: str = "none" +) -> mx.array: """ Computes the Kullback-Leibler divergence loss between logits and targets. @@ -98,7 +107,8 @@ def kl_div_loss(logits: mx.array, targets: mx.array, axis: int = -1, reduction: return _reduce(loss, reduction) -def _reduce(loss: mx.array, reduction: str = 'none'): + +def _reduce(loss: mx.array, reduction: str = "none"): if reduction == "mean": return mx.mean(loss) elif reduction == "sum": diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index 19ef2eddd..f9bb6a200 100644 --- a/python/tests/test_nn.py +++ b/python/tests/test_nn.py @@ -42,11 +42,11 @@ class TestNN(mlx_tests.MLXTestCase): targets = mx.array([0.5, 0.2, 0.9, 0.0]) losses = nn.losses.l1_loss(predictions, targets, reduction="none") self.assertEqual(losses, 0.0) - + def test_mse_loss(self): predictions = mx.array([0.5, 0.2, 0.9, 0.0]) targets = mx.array([0.7, 0.1, 0.8, 0.2]) - + # Test with reduction 'none' losses_none = nn.losses.mse_loss(predictions, targets, reduction="none") expected_none = mx.array([0.04, 0.01, 0.01, 0.04]) @@ -62,7 +62,6 @@ class TestNN(mlx_tests.MLXTestCase): expected_sum = mx.sum(expected_none) self.assertEqual(losses_sum, expected_sum) - def test_nll_loss(self): logits = mx.array([[0.0, -float("inf")], [-float("inf"), 0.0]]) targets = mx.array([0, 1]) @@ -82,7 +81,6 @@ class TestNN(mlx_tests.MLXTestCase): expected_sum = mx.sum(expected_none) self.assertEqual(losses_sum, expected_sum) - def test_kl_div_loss(self): p_logits = mx.array([[1.0, 2.0], [0.5, 1.0]]) q_logits = mx.array([[0.8, 1.5], [0.4, 1.2]])