| 46 | } |
| 47 | |
| 48 | void init_ops(nb::module_& m) { |
| 49 | m.def( |
| 50 | "reshape", |
| 51 | &mx::reshape, |
| 52 | nb::arg(), |
| 53 | "shape"_a, |
| 54 | nb::kw_only(), |
| 55 | "stream"_a = nb::none(), |
| 56 | nb::sig( |
| 57 | "def reshape(a: array, /, shape: Sequence[int], *, stream: " |
| 58 | "Union[None, Stream, Device] = None) -> array"), |
| 59 | R"pbdoc( |
| 60 | Reshape an array while preserving the size. |
| 61 | |
| 62 | Args: |
| 63 | a (array): Input array. |
| 64 | shape (tuple(int)): New shape. |
| 65 | stream (Stream, optional): Stream or device. Defaults to ``None`` |
| 66 | in which case the default stream of the default device is used. |
| 67 | |
| 68 | Returns: |
| 69 | array: The reshaped array. |
| 70 | )pbdoc"); |
| 71 | m.def( |
| 72 | "flatten", |
| 73 | [](const mx::array& a, |
| 74 | int start_axis, |
| 75 | int end_axis, |
| 76 | const mx::StreamOrDevice& s) { |
| 77 | return mx::flatten(a, start_axis, end_axis); |
| 78 | }, |
| 79 | nb::arg(), |
| 80 | "start_axis"_a = 0, |
| 81 | "end_axis"_a = -1, |
| 82 | nb::kw_only(), |
| 83 | "stream"_a = nb::none(), |
| 84 | nb::sig( |
| 85 | "def flatten(a: array, /, start_axis: int = 0, end_axis: int = " |
| 86 | "-1, *, stream: Union[None, Stream, Device] = None) -> array"), |
| 87 | R"pbdoc( |
| 88 | Flatten an array. |
| 89 | |
| 90 | The axes flattened will be between ``start_axis`` and ``end_axis``, |
| 91 | inclusive. Negative axes are supported. After converting negative axis to |
| 92 | positive, axes outside the valid range will be clamped to a valid value, |
| 93 | ``start_axis`` to ``0`` and ``end_axis`` to ``ndim - 1``. |
| 94 | |
| 95 | Args: |
| 96 | a (array): Input array. |
| 97 | start_axis (int, optional): The first dimension to flatten. Defaults to ``0``. |
| 98 | end_axis (int, optional): The last dimension to flatten. Defaults to ``-1``. |
| 99 | stream (Stream, optional): Stream or device. Defaults to ``None`` |
| 100 | in which case the default stream of the default device is used. |
| 101 | |
| 102 | Returns: |
| 103 | array: The flattened array. |
| 104 | |
| 105 | Example: |
no test coverage detected