RBERT (version 0.1.11)

transpose_for_scores: Reshape and transpose tensor

Description

In Python code, this is internal to attention_layer. Pulling it out into separate function here.

Usage

transpose_for_scores(
  input_tensor,
  batch_size,
  num_attention_heads,
  seq_length,
  width
)

Arguments

input_tensor

Tensor to reshape and transpose.

batch_size

Size of the first dimension of input_tensor.

num_attention_heads

Size of the third dimension of input_tensor. (Will be transposed to second dimension.)

seq_length

Size of the second dimension of input_tensor. (Will be transposed to third dimension.)

width

Size of fourth dimension of input_tensor.

Value

Tensor of shape: batch_size, num_attention_heads, seq_length, width.