diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp
index b2b7306dd..eb5d9d6b3 100644
--- a/mlx/primitives.cpp
+++ b/mlx/primitives.cpp
@@ -620,10 +620,11 @@ std::vector<array> ArgReduce::vjp(
 }
 
 std::vector<array> ArgReduce::jvp(
+    const std::vector<array>& primals,
     const std::vector<array>&,
-    const std::vector<array>& tangents,
     const std::vector<int>&) {
-  return {zeros_like(tangents[0], stream())};
+  auto shape = output_shapes(primals)[0];
+  return {zeros(shape, uint32, stream())};
 }
 
 std::pair<std::vector<array>, std::vector<int>> ArgSort::vmap(
@@ -647,6 +648,21 @@ bool ArgSort::is_equivalent(const Primitive& other) const {
   return axis_ == r_other.axis_;
 }
 
+std::vector<array> ArgSort::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>&,
+    const std::vector<int>&,
+    const std::vector<array>&) {
+  return {zeros_like(primals[0], stream())};
+}
+
+std::vector<array> ArgSort::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>&,
+    const std::vector<int>&) {
+  return {zeros(primals[0].shape(), uint32, stream())};
+}
+
 std::vector<array> AsType::vjp(
     const std::vector<array>& primals,
     const std::vector<array>& cotangents,
diff --git a/mlx/primitives.h b/mlx/primitives.h
index f4f157298..3d3202aaa 100644
--- a/mlx/primitives.h
+++ b/mlx/primitives.h
@@ -378,6 +378,7 @@ class ArgSort : public UnaryPrimitive {
   void eval_gpu(const std::vector<array>& inputs, array& out) override;
 
   DEFINE_VMAP()
+  DEFINE_GRADS()
   DEFINE_PRINT(ArgSort)
   DEFINE_INPUT_OUTPUT_SHAPE()
   bool is_equivalent(const Primitive& other) const override;