message: Support XLA DataFormatVecPermute
diff: @@ -91,15 +91,15 @@ class DataFormatVecPermuteOp : public XlaOpKernel {
  : XlaOpKernel(ctx) {
  OP_REQUIRES_OK(ctx, ctx->GetAttr("src_format", &src_format_));
  OP_REQUIRES(
  - ctx, src_format_.size() == 4,
  - errors::InvalidArgument("Data format should have 4 characters"));
  + ctx, src_format_.size() == 4 || src_format_.size() == 5,
  + errors::InvalidArgument("Data format should have 4 or 5 characters"));
  TensorFormat data_format;
  OP_REQUIRES(ctx, FormatFromString(src_format_, &data_format),
  errors::InvalidArgument("Invalid data format"));
  OP_REQUIRES_OK(ctx, ctx->GetAttr("dst_format", &dst_format_));
  OP_REQUIRES(
  - ctx, dst_format_.size() == 4,
  - errors::InvalidArgument("Data format should have 4 characters"));
  + ctx, dst_format_.size() == 4 || dst_format_.size() == 5,
  + errors::InvalidArgument("Data format should have 4 or 5 characters"));
  OP_REQUIRES(ctx, FormatFromString(dst_format_, &data_format),
  errors::InvalidArgument("Invalid data format"));
  }
  @@ -113,10 +113,10 @@ class DataFormatVecPermuteOp : public XlaOpKernel {
  input_tensor_shape.DebugString()));
  const int dim0 = input_tensor_shape.dim_size(0);
  OP_REQUIRES(
  - ctx, dim0 == 2 || dim0 == 4,
  + ctx, dim0 == 2 || dim0 == 4 || dim0 == 5,
  errors::InvalidArgument(
  - "First dimension of input must be of size 4, but got shape ",
  - input_tensor_shape.DebugString()));
  + "First dimension of input must be of size 2, 4 or 5, but got "
  + "shape ", input_tensor_shape.DebugString()));
  if (input_rank == 2) {
  OP_REQUIRES(
  ctx, input_tensor_shape.dim_size(1) == 2,
  