       let logp_lex = logp_lex_batch[pos].scalarsADHack  // [strVocab.chr.count]
       let logp_chr = decode(candidates, current_state).scalarsADHack  // [candidates.count]
       if pos != 0 {
-        // TODO: Mutate in place when AD supports it.
         let updatedNode = Lattice.Node(
           bestEdge: lattice[pos].bestEdge,
           bestScore: lattice[pos].bestScore,
           edges: lattice[pos].edges,
           semiringScore: lattice[pos].computeSemiringScore()
         )
-        lattice.positions = update(lattice.positions, index: pos, value: updatedNode)
+        lattice.positions.update(at: pos, to: updatedNode)
       }
 
       //for i, candidate in enumerate(candidates):

    return logp_lex[Int(index)]
  }

  // TODO: Triggers compiler crash.
  @differentiable
  public func buildLattice(_ sentence: CharacterSequence, maxLen: Int) -> Lattice {
    var lattice = Lattice(count: sentence.count, embEnc.embeddings)

          previous: lattice[pos].semiringScore,
          order: parameters.order)

        let updatedNode = Lattice.Node(
          bestEdge: lattice[next_pos].bestEdge,
          bestScore: lattice[next_pos].bestScore,
          edges: lattice[next_pos].edges + [edge],
          semiringScore: lattice[next_pos].semiringScore
        )
        lattice.positions.update(at: next_pos, to: updatedNode)
      }
    }

    // lattice[sentence.count].recomputeSemiringScore()
    let updatedNode = Lattice.Node(
      bestEdge: lattice[sentence.count].bestEdge,
      bestScore: lattice[sentence.count].bestScore,
      edges: lattice[sentence.count].edges,
      semiringScore: lattice[sentence.count].computeSemiringScore()
    )
    lattice.positions.update(at: sentence.count, to: updatedNode)

    return lattice
  }
}

extension Array {
  // NOTE(TF-1277): this mutating method exists as a workaround for `Array.subscript._modify` not
  // being differentiable.
  //
  // Semantically, it behaves like `Array.subscript.set`.
  @inlinable
  mutating func update(at index: Int, to value: Element) {
    self[index] = value
  }

  @usableFromInline
  @derivative(of: update)
  mutating func vjpUpdate(at index: Int, to value: Element) -> (
    value: (),
    pullback: (inout TangentVector) -> Element.TangentVector
  ) where Element: Differentiable {
    update(at: index, to: value)
    func pullback(_ dSelf: inout TangentVector) -> Element.TangentVector {
      let dElement = dSelf[index]
      dSelf.base[index] = dElement.zeroTangentVector
      return dElement
    }
    return ((), pullback)
  }
}

public struct MLP: Layer {
