diff --git a/docs/src/python/transforms.rst b/docs/src/python/transforms.rst
index 23f86720..566af4ac 100644
--- a/docs/src/python/transforms.rst
+++ b/docs/src/python/transforms.rst
@@ -11,6 +11,7 @@ Transforms
    eval
    async_eval
    compile
+   checkpoint
    custom_function
    disable_compile
    enable_compile
diff --git a/python/mlx/_stub_patterns.txt b/python/mlx/_stub_patterns.txt
index 637d7de4..7ca43698 100644
--- a/python/mlx/_stub_patterns.txt
+++ b/python/mlx/_stub_patterns.txt
@@ -1,10 +1,12 @@
 mlx.core.__prefix__:
-  from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
+  from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union, ParamSpec, TypeVar
   import sys
   if sys.version_info >= (3, 10):
     from typing import TypeAlias
   else:
     from typing_extensions import TypeAlias
+  P = ParamSpec("P")
+  R = TypeVar("R")
 
 mlx.core.__suffix__:
   from typing import Union
diff --git a/python/src/transforms.cpp b/python/src/transforms.cpp
index 67863038..c615b12e 100644
--- a/python/src/transforms.cpp
+++ b/python/src/transforms.cpp
@@ -1333,7 +1333,7 @@ void init_transforms(nb::module_& m) {
       "argnums"_a = nb::none(),
       "argnames"_a = std::vector<std::string>{},
       nb::sig(
-          "def value_and_grad(fun: Callable, argnums: Optional[Union[int, Sequence[int]]] = None, argnames: Union[str, Sequence[str]] = []) -> Callable"),
+          "def value_and_grad(fun: Callable[P, R], argnums: Optional[Union[int, Sequence[int]]] = None, argnames: Union[str, Sequence[str]] = []) -> Callable[P, Tuple[R, Any]]"),
       R"pbdoc(
         Returns a function which computes the value and gradient of ``fun``.
 
@@ -1402,7 +1402,7 @@ void init_transforms(nb::module_& m) {
       "argnums"_a = nb::none(),
       "argnames"_a = std::vector<std::string>{},
       nb::sig(
-          "def grad(fun: Callable, argnums: Optional[Union[int, Sequence[int]]] = None, argnames: Union[str, Sequence[str]] = []) -> Callable"),
+          "def grad(fun: Callable[P, R], argnums: Optional[Union[int, Sequence[int]]] = None, argnames: Union[str, Sequence[str]] = []) -> Callable[P, Any]"),
       R"pbdoc(
         Returns a function which computes the gradient of ``fun``.
 
@@ -1435,7 +1435,7 @@ void init_transforms(nb::module_& m) {
       "in_axes"_a = 0,
       "out_axes"_a = 0,
       nb::sig(
-          "def vmap(fun: Callable, in_axes: object = 0, out_axes: object = 0) -> Callable"),
+          "def vmap(fun: Callable[P, R], in_axes: object = 0, out_axes: object = 0) -> Callable[P, R]"),
       R"pbdoc(
         Returns a vectorized version of ``fun``.
 
@@ -1472,7 +1472,7 @@ void init_transforms(nb::module_& m) {
       "outputs"_a = nb::none(),
       "shapeless"_a = false,
       nb::sig(
-          "def compile(fun: Callable, inputs: Optional[object] = None, outputs: Optional[object] = None, shapeless: bool = False) -> Callable"),
+          "def compile(fun: Callable[P, R], inputs: Optional[object] = None, outputs: Optional[object] = None, shapeless: bool = False) -> Callable[P, R]"),
       R"pbdoc(
         Returns a compiled function which produces the same output as ``fun``.
 
@@ -1518,7 +1518,22 @@ void init_transforms(nb::module_& m) {
   m.def(
       "checkpoint",
       [](nb::callable fun) { return mlx_func(PyCheckpointedFun{fun}, fun); },
-      "fun"_a);
+      "fun"_a,
+      nb::sig("def checkpoint(fun: Callable[P, R]) -> Callable[P, R]"),
+      R"pbdoc(
+      Transform the passed callable to one that performs gradient
+      checkpointing with respect to the inputs of the callable.
+
+      Use this to reduce memory use for gradient computations at the expense of
+      increased computation.
+
+      Args:
+          fun (Callable): The function to checkpoint.
+
+      Returns:
+          A callable that recomputes intermediate states during gradient
+          computation.
+      )pbdoc");
 
   // Register static Python object cleanup before the interpreter exits
   auto atexit = nb::module_::import_("atexit");