26359ff3da29f088976a5bb4d6c67abe84b8b983,thinc/backends/jax_ops.py,,backprop_lstm_stepper,#,720

Before Change

def backprop_lstm_stepper(t, state):
    (dW, db, dX), (dY, dCt3), (G, C, S), (W, b) = state
    // TODO: I think I"m ignoring the sequence grads here?
    (dAt2, dYt2, dCt2) = backprop_lstm_gates(dCt3, dY[t+1], G[t], C[t+1], C[t])
    dX3, dY_prev, dW3, db3 = backprop_lstm_weights(dAt, (S[t], W, b))
    dX = index_update(dX, index[t], dX3)
    return (dW+dW3, db+db3, dX), (dY, dC), (G, C, X), (W, b)

After Change

    dAt3, dCt2 = backprop_lstm_gates(dCt3, dYt3, Gt3, Ct3, Ct2)
    dXt3, dYt2, dW3, db3 = backprop_lstm_weights(dAt3, (St3, W, b))
    dX = index_update(dX, index[t], dX3)
    dY = index_update(dY, index[t], dYt2)
    return (dW+dW3, db+db3, dX), (dY, dCt2), (G, C, X), (W, b)


@jax_jit()

Instances