(decoder_args, batchxbeam, memory_max_seq_len, decoding_max_seq_len)
| 302 | |
| 303 | |
| 304 | def init_op_cache(decoder_args, batchxbeam, memory_max_seq_len, decoding_max_seq_len): |
| 305 | use_batch_major_op_cache, x = get_op_cache_config(decoder_args.size_per_head, decoder_args.dtype) |
| 306 | if use_batch_major_op_cache == True: |
| 307 | self_cache = ( tf.zeros([decoder_args.num_layer, batchxbeam, decoder_args.head_num, decoder_args.size_per_head // x, |
| 308 | decoding_max_seq_len, x], dtype=decoder_args.dtype, name="op_self_cache_keys"), |
| 309 | tf.zeros([decoder_args.num_layer, batchxbeam, decoder_args.head_num, decoding_max_seq_len, |
| 310 | decoder_args.size_per_head], dtype=decoder_args.dtype, name="op_self_cache_values") ) |
| 311 | # use old format for now |
| 312 | # mem_cache = tf.zeros([decoder_args.num_layer, 2, batchxbeam, memory_max_seq_len, |
| 313 | # decoder_args.hidden_dim], dtype=decoder_args.dtype, name="op_memory_caches") |
| 314 | mem_cache = tf.zeros([2, decoder_args.num_layer, batchxbeam, |
| 315 | memory_max_seq_len, decoder_args.hidden_dim], dtype=decoder_args.dtype, name="op_memory_caches") |
| 316 | else : |
| 317 | # TODO(bhsueh) remove this case |
| 318 | self_cache = ( tf.zeros([decoder_args.num_layer, 0, batchxbeam, |
| 319 | decoder_args.hidden_dim], dtype=decoder_args.dtype, name="op_self_cache_keys"), |
| 320 | tf.zeros([decoder_args.num_layer, 0, batchxbeam, |
| 321 | decoder_args.hidden_dim], dtype=decoder_args.dtype, name="op_self_cache_values") ) |
| 322 | # mem_cache = tf.zeros([decoder_args.num_layer, 2, batchxbeam, |
| 323 | # memory_max_seq_len, decoder_args.hidden_dim], dtype=decoder_args.dtype, name="op_memory_caches") |
| 324 | mem_cache = tf.zeros([2, decoder_args.num_layer, batchxbeam, |
| 325 | memory_max_seq_len, decoder_args.hidden_dim], dtype=decoder_args.dtype, name="op_memory_caches") |
| 326 | |
| 327 | return self_cache, mem_cache |
| 328 | |
| 329 | def op_decoder(inputs, |
| 330 | memory_tensor, |
no test coverage detected