[WordSeg] Fix pullback for Array.update(at:to:). (#653)

The pullback of `Array.update(at:to:)` resets `dSelf.base[index]` to a zero value.

For correctness, it should use the value of `self[index].zeroTangentVector` before
the update occurs, not `dElement.zeroTangentVector`.

Author: dan-zheng

Models/Text/WordSeg/Model.swift

@@ -335,10 +335,11 @@ extension Array {
     value: (),
     pullback: (inout TangentVector) -> Element.TangentVector
   ) where Element: Differentiable {
+    let elementZero = self[index].zeroTangentVector
     update(at: index, to: value)
     func pullback(_ dSelf: inout TangentVector) -> Element.TangentVector {
       let dElement = dSelf[index]
-      dSelf.base[index] = dElement.zeroTangentVector
+      dSelf.base[index] = elementZero
       return dElement
     }
     return ((), pullback)