Module ops

Module ops 

Source
Expand description

TPU-optimized operations

Functionsยง

attention_tpu
Transformer attention (optimized for TPU)
batch_matmul_tpu
Batch matrix multiplication (optimized for TPU)
conv2d_tpu
Convolution on TPU
matmul_tpu
Matrix multiplication on TPU