(self, inputs, state)
Update input cache for fast inference.
Update input cache for fast inference.
Update input cache for fast inference .
def _use_predict_mem(self, inputs, state):
    """Update input cache for fast inference."""

    # inputs is (qk, v). mask isn't passed in.
    # where qk/v are shaped - (batch * n_heads, seq_len, d_head)

    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)
def _use_predict_mem ( self , inputs , state ) : # inputs is (qk, v). mask isn't passed in. # where qk/v are shaped - (batch * n_heads, seq_len, d_head) 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 )