| 78 | } // namespace |
| 79 | |
| 80 | void init_fast(nb::module_& parent_module) { |
| 81 | auto m = |
| 82 | parent_module.def_submodule("fast", "mlx.core.fast: fast operations"); |
| 83 | |
| 84 | m.def( |
| 85 | "rms_norm", |
| 86 | &mx::fast::rms_norm, |
| 87 | "x"_a, |
| 88 | "weight"_a.none(), |
| 89 | "eps"_a, |
| 90 | nb::kw_only(), |
| 91 | "stream"_a = nb::none(), |
| 92 | nb::sig( |
| 93 | "def rms_norm(x: array, weight: Optional[array], eps: float, *, stream: Union[None, Stream, Device] = None) -> array"), |
| 94 | R"pbdoc( |
| 95 | Root Mean Square normalization (RMS norm). |
| 96 | |
| 97 | The normalization is with respect to the last axis of the input ``x``. |
| 98 | |
| 99 | Args: |
| 100 | x (array): Input array. |
| 101 | weight (array, optional): A multiplicative weight to scale the result by. |
| 102 | The ``weight`` should be one-dimensional with the same size |
| 103 | as the last axis of ``x``. If set to ``None`` then no scaling happens. |
| 104 | eps (float): A small additive constant for numerical stability. |
| 105 | |
| 106 | Returns: |
| 107 | array: The output array. |
| 108 | )pbdoc"); |
| 109 | |
| 110 | m.def( |
| 111 | "layer_norm", |
| 112 | &mx::fast::layer_norm, |
| 113 | "x"_a, |
| 114 | "weight"_a.none(), |
| 115 | "bias"_a.none(), |
| 116 | "eps"_a, |
| 117 | nb::kw_only(), |
| 118 | "stream"_a = nb::none(), |
| 119 | nb::sig( |
| 120 | "def layer_norm(x: array, weight: Optional[array], bias: Optional[array], eps: float, *, stream: Union[None, Stream, Device] = None) -> array"), |
| 121 | R"pbdoc( |
| 122 | Layer normalization. |
| 123 | |
| 124 | The normalization is with respect to the last axis of the input ``x``. |
| 125 | |
| 126 | Args: |
| 127 | x (array): Input array. |
| 128 | weight (array, optional): A multiplicative weight to scale the result by. |
| 129 | The ``weight`` should be one-dimensional with the same size |
| 130 | as the last axis of ``x``. If set to ``None`` then no scaling happens. |
| 131 | bias (array, optional): An additive offset to be added to the result. |
| 132 | The ``bias`` should be one-dimensional with the same size |
| 133 | as the last axis of ``x``. If set to ``None`` then no translation happens. |
| 134 | eps (float): A small additive constant for numerical stability. |
| 135 | |
| 136 | Returns: |
| 137 | array: The output array. |
no test coverage detected