Bruce Fontaine e994fb8f32 Add option for input tensors for TPUEmbedding to have a first dimension which is a multiple of the batch_size. ...

PiperOrigin-RevId: 354344093
Change-Id: I0fbb6820b9fa0bb0128ed4a27e2f527937741cbf

2021-01-28 10:55:20 -08:00

..

tpu_ops.py

Add option for input tensors for TPUEmbedding to have a first dimension which is a multiple of the batch_size.

2021-01-28 10:55:20 -08:00

tpu_ordinal_selector_op.py

…