fix ad_coop_vec_pack() derivative

Author: wjakob

src/extra/autodiff.cpp

@@ -3808,7 +3808,7 @@ class CoopVecPack : public dr::detail::CustomOpBase {
 
     void forward() override {
         std::lock_guard<Lock> guard(state.lock);
-        uint32_t size = (uint32_t) m_input_indices.size();
+        uint32_t size = (uint32_t) m_inputs.size();
         JitIndex *tmp = (JitIndex *) alloca(sizeof(JitIndex) * size);
         size_t n_valid = 0;
 
@@ -3835,7 +3835,7 @@ class CoopVecPack : public dr::detail::CustomOpBase {
 
     void backward() override {
         std::lock_guard<Lock> guard(state.lock);
-        uint32_t n = (uint32_t) m_input_indices.size();
+        uint32_t n = (uint32_t) m_inputs.size();
 
         Variable *v = state[m_output_indices[0]];
         if (!v->grad.valid())
@@ -3844,9 +3844,13 @@ class CoopVecPack : public dr::detail::CustomOpBase {
         JitIndex *tmp = (JitIndex *) alloca(sizeof(JitIndex) * n);
         jit_coop_vec_unpack(v->grad.index(), n, tmp);
 
-        for (size_t i = 0; i < m_input_indices.size(); ++i) {
-            Variable *v2 = state[m_inputs[i]];
-            v2->accum(JitVar::steal(tmp[i]), v2->size);
+        for (size_t i = 0; i < m_inputs.size(); ++i) {
+            uint32_t index = m_inputs[i];
+            JitVar var = JitVar::steal(tmp[i]);
+            if (!index)
+                continue;
+            Variable *v2 = state[index];
+            v2->accum(var, v2->size);
         }
     }