def _use_predict_mem(self, inputs, state):
    """Update input cache for fast inference."""
    mem_end, mem, state = state
    seqlen = inputs[0].shape[-2]

    if seqlen <= self._predict_drop_len and seqlen < self._predict_mem_len:
      # This branch is called when only a small number of tokens are appended to
      # the sequence, e.g. when generating one token at a time. A fixed number
      # of tokens (self._predict_drop_tokens) will be dropped from memory if
      # needed, and then new values will be inserted into the memory.
      def roll_mem(buf):
        return np.concatenate(
            [buf[:, self._predict_drop_len:],
             np.zeros_like(buf[:, :self._predict_drop_len])], axis=1)

      do_roll_mem = (mem_end + seqlen > self._predict_mem_len)
      mem = fastmath.cond(
          pred=do_roll_mem,
          true_operand=mem,
          true_fun=lambda x: fastmath.nested_map(roll_mem, x),
          false_operand=mem,
          false_fun=lambda x: x,
      )
      mem_end = np.where(do_roll_mem, mem_end - self._predict_drop_len, mem_end)
      def update_mem(mem_element, new_vals):
        assert new_vals.shape[1] == seqlen
        if seqlen == 1:
          return fastmath.index_update(
              mem_element, jax.numpy.index_exp[:, mem_end], new_vals[:, 0, ...])
        else:
          return fastmath.dynamic_update_slice_in_dim(
              mem_element, new_vals, mem_end, axis=1)
      inputs = fastmath.nested_map_multiarg(update_mem, mem, inputs)
      return inputs, state, mem_end, inputs, mem_end + seqlen
    else:
      assert seqlen > self._predict_drop_len or seqlen == self._predict_mem_len
      # This branch handles the case where a large number of tokens are being
      # introduced all at once. The code here assumes that we are at the start
      # of the sequence, which matches the typical use case of decoding from a
      # language model given a long prefix. Note that if we're not at the start
      # of the sequence, the code here won't work.
      new_flat_mem = []
      for inp in fastmath.tree_leaves(inputs):
        assert inp.shape[1] == seqlen
        if seqlen == self._predict_mem_len:
          new_mem_val = inp
        elif seqlen > self._predict_mem_len:
          new_mem_val = inp[:, -self._predict_mem_len:]  # pylint: disable=invalid-unary-operand-type
        else:
          new_mem_val = np.concatenate([
              inp,
              np.zeros(inp.shape[:1]
                       + (self._predict_mem_len - inp.shape[1],)
                       + inp.shape[2:],
                       dtype=inp.dtype)
          ], axis=1)
        new_flat_mem.append(new_mem_val)
      mem, _ = fastmath.tree_unflatten(new_flat_mem, mem)

      # This code only works at the start of the sequence. There's no "assert"
      # primitive we can use to signal an error, so we instead signal the error
      # by introducing NaNs into the computation.
      def replace_with_nan_if_not_seq_start(x):
        if x.dtype != np.float32:
          return x
        return fastmath.cond(
            pred=np.equal(mem_end, np.array(0, dtype=mem_end.dtype)),
            true_operand=x, true_fun=lambda x: x,
            false_operand=x, false_fun=lambda x: x * np.nan)
      inputs = fastmath.nested_map(replace_with_nan_if_not_seq_start, inputs)
      return inputs, state, 0, mem, np.minimum(seqlen, self._predict_mem_len)
Update input cache for fast inference.
trax/layers/research/efficient_attention.py