| 16 | using namespace nb::literals; |
| 17 | |
| 18 | void init_linalg(nb::module_& parent_module) { |
| 19 | auto m = parent_module.def_submodule( |
| 20 | "linalg", "mlx.core.linalg: linear algebra routines."); |
| 21 | |
| 22 | m.def( |
| 23 | "norm", |
| 24 | [](const mx::array& a, |
| 25 | const std::variant<std::monostate, int, double, std::string>& ord_, |
| 26 | const std::variant<std::monostate, int, std::vector<int>>& axis_, |
| 27 | const bool keepdims, |
| 28 | const mx::StreamOrDevice stream) { |
| 29 | std::optional<std::vector<int>> axis = std::nullopt; |
| 30 | if (auto pv = std::get_if<int>(&axis_); pv) { |
| 31 | axis = std::vector<int>{*pv}; |
| 32 | } else if (auto pv = std::get_if<std::vector<int>>(&axis_); pv) { |
| 33 | axis = *pv; |
| 34 | } |
| 35 | |
| 36 | if (std::holds_alternative<std::monostate>(ord_)) { |
| 37 | return mx::linalg::norm(a, axis, keepdims, stream); |
| 38 | } else { |
| 39 | if (auto pv = std::get_if<std::string>(&ord_); pv) { |
| 40 | return mx::linalg::norm(a, *pv, axis, keepdims, stream); |
| 41 | } |
| 42 | double ord; |
| 43 | if (auto pv = std::get_if<int>(&ord_); pv) { |
| 44 | ord = *pv; |
| 45 | } else { |
| 46 | ord = std::get<double>(ord_); |
| 47 | } |
| 48 | return mx::linalg::norm(a, ord, axis, keepdims, stream); |
| 49 | } |
| 50 | }, |
| 51 | nb::arg(), |
| 52 | "ord"_a = nb::none(), |
| 53 | "axis"_a = nb::none(), |
| 54 | "keepdims"_a = false, |
| 55 | nb::kw_only(), |
| 56 | "stream"_a = nb::none(), |
| 57 | nb::sig( |
| 58 | "def norm(a: array, /, ord: Union[None, int, float, str] = None, axis: Union[None, int, list[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), |
| 59 | R"pbdoc( |
| 60 | Matrix or vector norm. |
| 61 | |
| 62 | This function computes vector or matrix norms depending on the value of |
| 63 | the ``ord`` and ``axis`` parameters. |
| 64 | |
| 65 | Args: |
| 66 | a (array): Input array. If ``axis`` is ``None``, ``a`` must be 1-D or 2-D, |
| 67 | unless ``ord`` is ``None``. If both ``axis`` and ``ord`` are ``None``, the |
| 68 | 2-norm of ``a.flatten`` will be returned. |
| 69 | ord (int, float or str, optional): Order of the norm (see table under ``Notes``). |
| 70 | If ``None``, the 2-norm (or Frobenius norm for matrices) will be computed |
| 71 | along the given ``axis``. Default: ``None``. |
| 72 | axis (int or list(int), optional): If ``axis`` is an integer, it specifies the |
| 73 | axis of ``a`` along which to compute the vector norms. If ``axis`` is a |
| 74 | 2-tuple, it specifies the axes that hold 2-D matrices, and the matrix |
| 75 | norms of these matrices are computed. If `axis` is ``None`` then |