"tf.While"(%cst_0, %cst, %cst_0, %arg0, %cst_1) {T = [i32, i32, i32, f32, f32],_lower_using_switch_merge = true, _num_original_outputs = 5 : i64, _read_only_resource_inputs = [], body = @while_body, cond = @while_cond, device = "", is_stateless = true, output_shapes = [#tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<1x1024>, #tf_type.shape<1024x1024>], parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<i32>, tensor<1x1024xf32>, tensor<1024x1024xf32>) ->...

        init_func = @__init_func_530, next_func = @__next_func_680,
        next_func.experimental_ints_on_device = true,
        operandSegmentSizes = array<i32: 2, 2, 1>,
        output_shapes = [#tf_type.shape<?>],
        output_types = [f32],
        metadata = ""
      } :

        init_func = @__init_func_530, next_func = @__next_func_680,
        next_func.experimental_ints_on_device = true,
        operandSegmentSizes = array<i32: 2, 2, 1>,
        output_shapes = [#tf_type.shape<?>],
        output_types = [f32],
        metadata = ""
      } :

    %0:3 = "tf.While"(%cst_0, %cst_1, %arg0) {T = [i32, i32, f32], _lower_using_switch_merge = true, _num_original_outputs = 4 : i64, _read_only_resource_inputs = [], body = @while_body, cond = @while_cond, device = "", is_stateless = true, output_shapes = [#tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<1x5x5x1024>], parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<1x5x5x1024xf32>) -> (tensor<i32>, tensor<i32>, tensor<1x5x5x1024xf32>)

    "tf.Yield"(%1) : (tensor<5xf32>) -> ()
  }) {device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0", metadata = "", operandSegmentSizes = array<i32: 1, 2, 1>, output_shapes = [#tf_type.shape<>], output_types = [!tf_type.string]} : (tensor<4xf32>, tensor<3xf32>, tensor<!tf_type.resource>, tensor<2xf32>) -> tensor<!tf_type.variant>
  return
}

// -----

  SmallVector<int64_t, 2> output_shape{1, mlir::ShapedType::kDynamic};
  EXPECT_EQ(mlir::cast<RankedTensorType>(output_types[0]).getShape().size(),
            output_shape.size());
  for (int i = 0; i < output_shape.size(); i++) {
    EXPECT_EQ(mlir::cast<RankedTensorType>(output_types[0]).getDimSize(i),
              output_shape[i]);
  }
}

            s: "loc:@a"
          }
        }
      }
      attr {
        key: "_output_shapes"
        value {
          list {
            shape {
            }
          }
        }
      }
    }
    node {
      name: "b/initial_value"
      op: "Const"
      attr {
        key: "_output_shapes"
        value {
          list {
            shape {
            }
          }

  PartialTensorShape output_shape =
      ConvertTypeToTensorShape(mlir::UnrankedTensorType::get(b.getF32Type()));
  EXPECT_TRUE(output_shape.IsIdenticalTo(PartialTensorShape()));
}

TEST(ConvertTypeToTensorTypeTest, NonFullyDefinedRankedTensorType) {
  mlir::MLIRContext context;
  RegisterDialects(context);
  mlir::Builder b(&context);

  PartialTensorShape output_shape = ConvertTypeToTensorShape(

    return failure();

  // Build the lstm op.
  SmallVector<int64_t, 3> output_shape;
  if (time_majored) {
    output_shape = {time, batch, n_output};
  } else {
    output_shape = {batch, time, n_output};
  }
  auto result_type = mlir::RankedTensorType::get(
      output_shape,
      mlir::cast<RankedTensorType>(final_inputs.getType()).getElementType());

    SmallVector<int64_t, 4> output_shape;
    for (int i = 0; i < num_dimensions; ++i) {
      perm.push_back(perm_tensor.getValues<IntegerAttr>()[i].getInt());
      output_shape.push_back(input_shape[perm[i]]);

      // Check that the derived output shape matches the static shape.
      assert(!output_type.hasStaticShape() ||
             output_type.getShape()[i] == output_shape[i]);
    }