message: Avoid excessive cpu<->gpu memory swaps, compute shape ops on the CPU. This
  results in +10% perf improvement for tensor2tensor Transformer model training
  step times, and +37% perf improvement for tensor2tensor Transformer model
  decoding.
diff: @@ -2903,6 +2903,8 @@ def tensordot(a, b, axes, name=None):
  free_dims_static = None
  shape_a = array_ops.shape(a)
  rank_a = array_ops.rank(a)
  + # TODO(b/115583659): Automate this.
  + with ops.device("/cpu:0"):
  axes = ops.convert_to_tensor(axes, dtype=dtypes.int32, name="axes")
  axes = cast(axes >= 0, dtypes.int32) * axes + cast(
  axes < 0, dtypes.int32) * (
  