Function init_transforms

python/src/transforms.cpp:1022–1561

    {0, 0}};

void init_transforms(nb::module_& m) {
  nb::class_<PyCustomFunction>(
      m,
      "custom_function",
      nb::type_slots(py_custom_function_slots),
      R"pbdoc(
      Set up a function for custom gradient and vmap definitions.

      This class is meant to be used as a function decorator. Instances are
      callables that behave identically to the wrapped function. However, when
      a function transformation is used (e.g. computing gradients using
      :func:`value_and_grad`) then the functions defined via
      :meth:`custom_function.vjp`, :meth:`custom_function.jvp` and
      :meth:`custom_function.vmap` are used instead of the default transformation.

      Note, all custom transformations are optional. Undefined transformations
      fall back to the default behaviour.

      Example:

        .. code-block:: python

            import mlx.core as mx

            @mx.custom_function
            def f(x, y):
                return mx.sin(x) * y

            @f.vjp
            def f_vjp(primals, cotangent, output):
                x, y = primals
                return cotan * mx.cos(x) * y, cotan * mx.sin(x)

            @f.jvp
            def f_jvp(primals, tangents):
              x, y = primals
              dx, dy = tangents
              return dx * mx.cos(x) * y + dy * mx.sin(x)

            @f.vmap
            def f_vmap(inputs, axes):
              x, y = inputs
              ax, ay = axes
              if ay != ax and ax is not None:
                  y = y.swapaxes(ay, ax)
              return mx.sin(x) * y, (ax or ay)

      All ``custom_function`` instances behave as pure functions. Namely, any
      variables captured will be treated as constants and no gradients will be
      computed with respect to the captured arrays. For instance:

        .. code-block:: python

          import mlx.core as mx

          def g(x, y):
            @mx.custom_function
            def f(x):