code: def top_1_tpu(inputs):
    
    inputs_max = tf.reduce_max(inputs, axis=-1, keepdims=True)
    mask = tf.to_int32(tf.equal(inputs_max, inputs))
    index = tf.range(tf.shape(inputs)[-1]) * mask
    return tf.squeeze(inputs_max, -1), tf.reduce_max(index, axis=-1)
docs: find max and argmax over the last dimension.
  
    Works well on TPU
  
    Args:
      inputs: A tensor with shape [..., depth]
  
    Returns:
      values: a Tensor with shape [...]
      indices: a Tensor with shape [...]