From 38694a1b6ea434d80583f62c727ecb452d7ff7e9 Mon Sep 17 00:00:00 2001
From: lzhengning <lzhengning@gmail.com>
Date: Sun, 21 Nov 2021 16:50:45 +0800
Subject: [PATCH] fix: wrong grad of varslice like [::-1]

---
 python/jittor/__init__.py           |  2 +-
 python/jittor/src/ops/getitem_op.cc |  4 ++--
 python/jittor/src/ops/setitem_op.cc | 20 ++++++++++----------
 python/jittor/src/var_slices.h      | 11 +++++++++--
 python/jittor/test/test_rnn.py      | 13 +++++++++++++
 python/jittor/test/test_setitem.py  |  8 +++++++-
 6 files changed, 42 insertions(+), 16 deletions(-)

diff --git a/python/jittor/__init__.py b/python/jittor/__init__.py
index 38683e15..7907e3a8 100644
--- a/python/jittor/__init__.py
+++ b/python/jittor/__init__.py
@@ -9,7 +9,7 @@
 # file 'LICENSE.txt', which is part of this source code package.
 # ***************************************************************
 
-__version__ = '1.3.1.18'
+__version__ = '1.3.1.19'
 from jittor_utils import lock
 with lock.lock_scope():
     ori_int = int
diff --git a/python/jittor/src/ops/getitem_op.cc b/python/jittor/src/ops/getitem_op.cc
index e33b9d4c..5c761a79 100644
--- a/python/jittor/src/ops/getitem_op.cc
+++ b/python/jittor/src/ops/getitem_op.cc
@@ -391,9 +391,9 @@ VarPtr GetitemOp::grad(Var* out, Var* dout, Var* v, int v_index) {
     // need analysis the overlap attr os var slices
     for (int i=0; i<vs.n; i++)
         if (vs.slices[i].is_var()) {
-            return make_setitem(zeros, VarSlices(vs), dout, ns_add);
+            return make_setitem(zeros, VarSlices(vs, true), dout, ns_add);
         }
-    return make_setitem(zeros, VarSlices(vs), dout, ns_void);
+    return make_setitem(zeros, VarSlices(vs, true), dout, ns_void);
 }
 
 void GetitemOp::jit_prepare(JK& jk) {
diff --git a/python/jittor/src/ops/setitem_op.cc b/python/jittor/src/ops/setitem_op.cc
index 23f600ab..13836f36 100644
--- a/python/jittor/src/ops/setitem_op.cc
+++ b/python/jittor/src/ops/setitem_op.cc
@@ -128,41 +128,41 @@ VarPtr SetitemOp::grad(Var* out, Var* dout, Var* v, int v_index) {
         if (v_index == 0) {
             float32 number = 0;
             VarPtr zero = make_array(&number, 1, ns_float32);
-            return make_setitem(dout, VarSlices(vs), zero, ns_void);
+            return make_setitem(dout, VarSlices(vs, true), zero, ns_void);
         } else {
-            return make_getitem(dout, VarSlices(vs));
+            return make_getitem(dout, VarSlices(vs, true));
         }
     }
     if (op == ns_add) {
         if (v_index == 0) {
             return dout;
         } else {
-            return make_getitem(dout, VarSlices(vs));
+            return make_getitem(dout, VarSlices(vs, true));
         }
     }
     if (op == ns_subtract) {
         if (v_index == 0) {
             return dout;
         } else {
-            return make_unary(make_getitem(dout, VarSlices(vs)), ns_negative);
+            return make_unary(make_getitem(dout, VarSlices(vs, true)), ns_negative);
         }
     }
     if (op == ns_multiply) {
         if (v_index == 0) {
-            return make_setitem(dout, VarSlices(vs), input(1), ns_multiply);
+            return make_setitem(dout, VarSlices(vs, true), input(1), ns_multiply);
         } else {
             return make_binary(
-                make_getitem(inputs().front(), VarSlices(vs)),
-                make_getitem(dout, VarSlices(vs)), ns_multiply);
+                make_getitem(inputs().front(), VarSlices(vs, true)),
+                make_getitem(dout, VarSlices(vs, true)), ns_multiply);
         }
     }
     if (op == ns_divide) {
         if (v_index == 0) {
-            return make_setitem(dout, VarSlices(vs), input(1), ns_divide);
+            return make_setitem(dout, VarSlices(vs, true), input(1), ns_divide);
         } else {
             // dy = -dz*x / y^2
-            auto dout2 = make_getitem(dout, VarSlices(vs));
-            auto x = make_getitem(inputs().front(), VarSlices(vs));
+            auto dout2 = make_getitem(dout, VarSlices(vs, true));
+            auto x = make_getitem(inputs().front(), VarSlices(vs, true));
             auto y = v;
             auto ndz = make_unary(dout2, ns_negative);
             auto ndzx = make_binary(ndz, x, ns_multiply);
diff --git a/python/jittor/src/var_slices.h b/python/jittor/src/var_slices.h
index 235c03e2..95c4cf01 100644
--- a/python/jittor/src/var_slices.h
+++ b/python/jittor/src/var_slices.h
@@ -47,9 +47,16 @@ struct VarSlices  {
     inline VarSlices(VarSlices&& other) : slices(other.slices), n(other.n) {
         other.slices = nullptr;
     }
-    inline VarSlices(const VarSlices& other) : slices(new VarSlice[other.n]), n(other.n) {
-        for (int i=0; i<n; i++)
+    inline VarSlices(const VarSlices& other, bool negtive_set_none=false) : slices(new VarSlice[other.n]), n(other.n) {
+        for (int i=0; i<n; i++) {
             slices[i] = other.slices[i];
+            if (negtive_set_none && 
+                slices[i].is_slice() && 
+                slices[i].slice.step < 0 &&
+                slices[i].slice.stop < 0) {
+                slices[i].slice.mask |= 2;
+            }
+        }
     }
     inline void operator=(VarSlices&& other) {
         if (slices) delete[] slices;
diff --git a/python/jittor/test/test_rnn.py b/python/jittor/test/test_rnn.py
index 9cc37fb6..7795f8c4 100644
--- a/python/jittor/test/test_rnn.py
+++ b/python/jittor/test/test_rnn.py
@@ -327,6 +327,19 @@ class TestRNN(unittest.TestCase):
         np.testing.assert_allclose(th.detach().cpu().numpy(), jh.data, rtol=1e-03, atol=1e-06)
         np.testing.assert_allclose(tc.detach().cpu().numpy(), jc.data, rtol=1e-03, atol=1e-06)
 
+    def test_twobilinear_lstm(self):
+        x = jt.rand(5, 4, 10)
+        rnn1 = nn.LSTM(10, 20, bidirectional=True)
+        out1, _ = rnn1(x)
+        rnn2 = nn.LSTM(40, 20, bidirectional=True)
+        out2, _ = rnn2(out1)
+        target = jt.zeros_like(out2)
+        loss = nn.mse_loss(out2, target)
+
+        from jittor import optim
+        optimizer = optim.RMSprop(rnn1.parameters())
+        optimizer.step(loss)
+
     @skipIf(not jt.has_cuda, "No Cuda found")
     @jt.flag_scope(use_cuda=1)
     def test_cudnn_rnn(self):
diff --git a/python/jittor/test/test_setitem.py b/python/jittor/test/test_setitem.py
index b8320ae8..8bbf5e7f 100644
--- a/python/jittor/test/test_setitem.py
+++ b/python/jittor/test/test_setitem.py
@@ -225,8 +225,14 @@ class TestSetitem(unittest.TestCase):
         b = a[...,:,None,:2]
         assert b.shape == [2,4,1,2]
         np.testing.assert_allclose(b.data, a.data[...,:,None,:2])
-        
 
+    def test_flip_grad(self):
+        a = jt.rand(10)
+        b = a[::-1]
+        c = b[::-1]
+        d = c.sum()
+        jt.grad(d, [a])
+        
 
 if __name__ == "__main__":
     unittest.main()
\ No newline at end of file