Adds radians and degrees (#1011)

This commit is contained in:
Aneesh Shetty 2024-04-22 13:17:49 -05:00 committed by GitHub
parent 3d405fb3b1
commit d0dbfe0b97
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 120 additions and 1 deletions

View File

@ -43,6 +43,7 @@ Operations
cummin cummin
cumprod cumprod
cumsum cumsum
degrees
dequantize dequantize
diag diag
diagonal diagonal
@ -99,6 +100,7 @@ Operations
prod prod
quantize quantize
quantized_matmul quantized_matmul
radians
reciprocal reciprocal
repeat repeat
reshape reshape

View File

@ -2192,6 +2192,16 @@ array arctanh(const array& a, StreamOrDevice s /* = {} */) {
a.shape(), dtype, std::make_shared<ArcTanh>(to_stream(s)), {input}); a.shape(), dtype, std::make_shared<ArcTanh>(to_stream(s)), {input});
} }
array degrees(const array& a, StreamOrDevice s /* = {} */) {
auto dtype = at_least_float(a.dtype());
return multiply(a, array(180.0 / M_PI, dtype), s);
}
array radians(const array& a, StreamOrDevice s /* = {} */) {
auto dtype = at_least_float(a.dtype());
return multiply(a, array(M_PI / 180.0, dtype), s);
}
array log(const array& a, StreamOrDevice s /* = {} */) { array log(const array& a, StreamOrDevice s /* = {} */) {
auto dtype = at_least_float(a.dtype()); auto dtype = at_least_float(a.dtype());
auto input = astype(a, dtype, s); auto input = astype(a, dtype, s);

View File

@ -851,6 +851,12 @@ array arccosh(const array& a, StreamOrDevice s = {});
/** Inverse Hyperbolic Tangent of the elements of an array */ /** Inverse Hyperbolic Tangent of the elements of an array */
array arctanh(const array& a, StreamOrDevice s = {}); array arctanh(const array& a, StreamOrDevice s = {});
/** Convert the elements of an array from Radians to Degrees **/
array degrees(const array& a, StreamOrDevice s = {});
/** Convert the elements of an array from Degrees to Radians **/
array radians(const array& a, StreamOrDevice s = {});
/** Natural logarithm of the elements of an array. */ /** Natural logarithm of the elements of an array. */
array log(const array& a, StreamOrDevice s = {}); array log(const array& a, StreamOrDevice s = {});

View File

@ -1032,6 +1032,40 @@ void init_ops(nb::module_& m) {
Returns: Returns:
array: The inverse hyperbolic tangent of ``a``. array: The inverse hyperbolic tangent of ``a``.
)pbdoc"); )pbdoc");
m.def(
"degrees",
&mlx::core::degrees,
nb::arg(),
nb::kw_only(),
"stream"_a = nb::none(),
nb::sig(
"def degrees(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
R"pbdoc(
Convert angles from radians to degrees.
Args:
a (array): Input array.
Returns:
array: The angles in degrees.
)pbdoc");
m.def(
"radians",
&mlx::core::radians,
nb::arg(),
nb::kw_only(),
"stream"_a = nb::none(),
nb::sig(
"def radians(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
R"pbdoc(
Convert angles from degrees to radians.
Args:
a (array): Input array.
Returns:
array: The angles in radians.
)pbdoc");
m.def( m.def(
"log", "log",
&mlx::core::log, &mlx::core::log,

View File

@ -893,6 +893,22 @@ class TestOps(mlx_tests.MLXTestCase):
self.assertTrue(np.allclose(result, expected)) self.assertTrue(np.allclose(result, expected))
def test_degrees(self):
a = mx.array(
[0, math.pi / 4, math.pi / 2, math.pi, 3 * math.pi / 4, 2 * math.pi]
)
result = mx.degrees(a)
expected = np.degrees(a, dtype=np.float32)
self.assertTrue(np.allclose(result, expected))
def test_radians(self):
a = mx.array([0.0, 45.0, 90.0, 180.0, 270.0, 360.0])
result = mx.radians(a)
expected = np.radians(a, dtype=np.float32)
self.assertTrue(np.allclose(result, expected))
def test_log1p(self): def test_log1p(self):
a = mx.array([1, 0.5, 10, 100]) a = mx.array([1, 0.5, 10, 100])
result = mx.log1p(a) result = mx.log1p(a)

View File

@ -49,8 +49,9 @@ class TestVmap(mlx_tests.MLXTestCase):
"sin", "sin",
"sqrt", "sqrt",
"square", "square",
"degrees",
"radians",
] ]
ops = ["erfinv"]
for opname in ops: for opname in ops:
with self.subTest(op=opname): with self.subTest(op=opname):
op = getattr(mx, opname) op = getattr(mx, opname)

View File

@ -1154,6 +1154,56 @@ TEST_CASE("test arithmetic unary ops") {
CHECK(allclose(cos(x), expected).item<bool>()); CHECK(allclose(cos(x), expected).item<bool>());
} }
// Test degrees
{
array x(0.0);
CHECK_EQ(degrees(x).item<float>(), 0.0);
x = array(M_PI_2);
CHECK(degrees(x).item<float>() == doctest::Approx(90.0));
CHECK(array_equal(degrees(array({})), array({})).item<bool>());
// Integer input type
x = array(0);
CHECK_EQ(x.dtype(), int32);
CHECK_EQ(degrees(x).item<float>(), 0.0);
// Input is irregularly strided
x = broadcast_to(array(M_PI_2), {2, 2, 2});
CHECK(allclose(degrees(x), full({2, 2, 2}, 90.0)).item<bool>());
float angles[] = {0.0f, M_PI_2, M_PI, 3.0f * M_PI_2};
x = split(array(angles, {2, 2}), 2, 1)[0];
auto expected = array({0.0f, 180.0f}, {2, 1});
CHECK(allclose(degrees(x), expected).item<bool>());
}
// Test radians
{
array x(0.0);
CHECK_EQ(radians(x).item<float>(), 0.0);
x = array(90.0);
CHECK(radians(x).item<float>() == doctest::Approx(M_PI_2));
CHECK(array_equal(radians(array({})), array({})).item<bool>());
// Integer input type
x = array(90);
CHECK_EQ(x.dtype(), int32);
CHECK(radians(x).item<float>() == doctest::Approx(M_PI_2));
// Input is irregularly strided
x = broadcast_to(array(90.0f), {2, 2, 2});
CHECK(allclose(radians(x), full({2, 2, 2}, M_PI_2)).item<bool>());
x = split(array({0.0f, 90.0f, 180.0f, 270.0f}, {2, 2}), 2, 1)[0];
float angles[] = {0.0f, M_PI};
auto expected = array(angles, {2, 1});
CHECK(allclose(radians(x), expected).item<bool>());
}
// Test log // Test log
{ {
array x(0.0); array x(0.0);