mirror of
https://github.com/ml-explore/mlx.git
synced 2025-10-19 08:38:09 +08:00
Fix divide types + floor divide (//) (#138)
* divide types * fix black + test
This commit is contained in:
@@ -2,7 +2,7 @@
|
||||
|
||||
import os
|
||||
import unittest
|
||||
from typing import Callable, List, Tuple
|
||||
from typing import Callable, List, Tuple, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import numpy as np
|
||||
@@ -21,7 +21,7 @@ class MLXTestCase(unittest.TestCase):
|
||||
|
||||
def assertEqualArray(
|
||||
self,
|
||||
args: List[mx.array | float | int],
|
||||
args: List[Union[mx.array, float, int]],
|
||||
mlx_func: Callable[..., mx.array],
|
||||
expected: mx.array,
|
||||
atol=1e-2,
|
||||
|
@@ -236,6 +236,24 @@ class TestOps(mlx_tests.MLXTestCase):
|
||||
self.assertEqual(z.dtype, mx.float32)
|
||||
self.assertEqual(z.item(), 0.5)
|
||||
|
||||
x = x.astype(mx.float16)
|
||||
z = x / 4.0
|
||||
self.assertEqual(z.dtype, mx.float16)
|
||||
|
||||
x = x.astype(mx.float16)
|
||||
z = 4.0 / x
|
||||
self.assertEqual(z.dtype, mx.float16)
|
||||
|
||||
x = mx.array(5)
|
||||
y = mx.array(2)
|
||||
z = x / y
|
||||
self.assertEqual(z.dtype, mx.float32)
|
||||
self.assertEqual(z.item(), 2.5)
|
||||
|
||||
z = x // y
|
||||
self.assertEqual(z.dtype, mx.int32)
|
||||
self.assertEqual(z.item(), 2)
|
||||
|
||||
def test_remainder(self):
|
||||
for dt in [mx.int32, mx.float32]:
|
||||
x = mx.array(2, dtype=dt)
|
||||
|
Reference in New Issue
Block a user