mirror of
https://github.com/ml-explore/mlx.git
synced 2025-10-19 00:04:41 +08:00
Added ArcTan2 operation (#1079)
* Added ArcTan2 operation * Cleanup, bug fixes from code review * Minor cleanup, fixed Linux tests
This commit is contained in:
@@ -930,6 +930,25 @@ void init_ops(nb::module_& m) {
|
||||
Returns:
|
||||
array: The inverse tangent of ``a``.
|
||||
)pbdoc");
|
||||
m.def(
|
||||
"arctan2",
|
||||
&mlx::core::arctan2,
|
||||
nb::arg(),
|
||||
nb::arg(),
|
||||
nb::kw_only(),
|
||||
"stream"_a = nb::none(),
|
||||
nb::sig(
|
||||
"def arctan2(a: array, b: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
|
||||
R"pbdoc(
|
||||
Element-wise inverse tangent of the ratio of two arrays.
|
||||
|
||||
Args:
|
||||
a (array): Input array.
|
||||
b (array): Input array.
|
||||
|
||||
Returns:
|
||||
array: The inverse tangent of the ratio of ``a`` and ``b``.
|
||||
)pbdoc");
|
||||
m.def(
|
||||
"sinh",
|
||||
&mlx::core::sinh,
|
||||
|
@@ -1273,6 +1273,7 @@ class TestOps(mlx_tests.MLXTestCase):
|
||||
"arcsin": lambda primal, cotan: cotan / np.sqrt(1.0 - primal**2),
|
||||
"arccos": lambda primal, cotan: -cotan / np.sqrt(1.0 - primal**2),
|
||||
"arctan": lambda primal, cotan: cotan / (1.0 + primal**2),
|
||||
"arctan2": lambda primal, cotan: cotan / (1.0 + primal**2),
|
||||
"arcsinh": lambda primal, cotan: cotan / np.sqrt(primal**2 + 1),
|
||||
"arccosh": lambda primal, cotan: cotan / np.sqrt(primal**2 - 1),
|
||||
"arctanh": lambda primal, cotan: cotan / (1.0 - primal**2),
|
||||
|
Reference in New Issue
Block a user