Add amax gradient fn

facebookresearch · Dec 5, 2022 · 31196c0 · 31196c0
1 parent 027bac0
commit 31196c0
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/shumai/tensor/register_gradients.ts b/shumai/tensor/register_gradients.ts
@@ -48,6 +48,9 @@ const impls = {
   add: (ctx: GradContext) => {
     return possiblyReduce(ctx.backward_input, ctx)
   },
+  amax: (ctx: GradContext) => {
+    return ctx.backward_input
+  },
   conv2d: (ctx: GradContext) => {
     const [x, w, sx, sy, px, py, dx, dy, g] = <
       [Tensor, Tensor, number, number, number, number, number, number, number]