inline DenseElementsAttr GetShape(Value output_val, bool truncate = false) {
  auto output_shape = output_val.getType().dyn_cast<ShapedType>().getShape();

  SmallVector<int32_t> shape;
  shape.reserve(output_shape.size());

  bool needs_truncation = true;
  for (size_t dim_idx = 0; dim_idx < output_shape.size(); ++dim_idx) {
    int64_t dim = output_shape[dim_idx];
    if (truncate && needs_truncation && dim == 1) {
      continue;

  // dimensions.
  SmallVector<int64_t> output_shape(output_tensor_rank);
  for (int i = 0; i < output_tensor_rank; i++) {
    if (collapsed_dims.contains(i)) {
      // The collapsed dimension's size should have been 1, so it restores the
      // dimension with size 1.
      output_shape[i] = 1;
    } else {
      output_shape[i] = *shape_itr;
      shape_itr++;
    }
  }

    std::vector<PartialTensorShape> output_shapes;
    std::transform(iter.second.begin(), iter.second.end(),
                   std::back_inserter(output_shapes),
                   [](const InferredShape& inferred_shape) {
                     return inferred_shape.shape;
                   });
    Node* n = node_name_index[iter.first];
    n->AddAttr(kXlaInferredShapesAttrName, output_shapes);
  }

  return absl::OkStatus();

  SmallVector<int64_t> output_shape(input_shape.getShape().begin(),
                                    input_shape.getShape().end());
  for (int i : spatial_dims) {
    output_shape[i] += padding_values[2 * i] + padding_values[2 * i + 1];
  }

  return builder.create<TF::PadV2Op>(
      loc, RankedTensorType::get(output_shape, builder.getI8Type()), input,
      temp_padding,

    const mlir::TensorType cluster_func_output_type,
    const xla::OpSharding& output_sharding,
    mlir::Type* tiled_logical_computation_type) {
  const auto output_shape = cluster_func_output_type.getShape();
  auto new_output_shape = llvm::to_vector<4>(output_shape);
  auto dimension_to_splits_map =
      GetDimensionIndicesAndNumSplitsFromSharding(output_sharding);
  if (!dimension_to_splits_map.ok()) {

  }

  // Gather shapes for output.
  for (auto v : ddn.lhs_batch_dimensions()) {
    output_shape.push_back(lhs_shape[v]);
  }

  // Batch dimension is gathered from the right side.
  if (output_shape.empty()) {
    for (auto v : ddn.rhs_batch_dimensions()) {
      output_shape.push_back(rhs_shape[v]);
    }
  }

  // Gather remaining dimensions.

  %2 = "tf.Const"() {__op_key = 2: i32, device = "/device:CPU:0", value = dense<1> : tensor<i64>} : () -> tensor<i64>
  %3 = "tf.RangeDataset"(%0, %1, %2) {__op_key = 3: i32, device = "/device:CPU:0", output_shapes = [#tf_type.shape<>], output_types = [i64], metadata = ""} : (tensor<i64>, tensor<i64>, tensor<i64>) -> tensor<!tf_type.variant>
  // CHECK: tf_mlrt.executeop{{.*}}op: \22FlatMapDataset\22

      continue;
    }
    shape.dims = new int64_t[shape.num_dims];
    for (size_t j = 0; j < shape.num_dims; ++j) {
      shape.dims[j] = c.Value(c.Dim(shape_handle, j));
    }
  }
  if (output_shapes != nullptr) *output_shapes = output_shapes_result;

  // TODO(bgogul): Set output_resource_shapes_and_types.
}

void TF_ImportGraphDefOptionsSetValidateColocationConstraints(

    const xla::HloInputOutputAliasConfig& input_output_alias,
    absl::Span<const int> input_mapping,
    const std::map<int, const Tensor*>& resource_vars_snapshots,
    DataType output_dtype, const TensorShape& output_shape,
    Allocator* output_allocator, bool allocate_xla_tensors, se::Stream* stream,
    bool use_multiple_streams, std::shared_ptr<se::Event> definition_event) {

  return ConvertTensor(t, builder);
}

void ConvertToTensorShapeProto(ArrayRef<int64_t> shape,
                               TensorShapeProto* output_shape) {
  for (auto d : shape) {
    output_shape->add_dim()->set_size(ShapedType::isDynamic(d) ? kTFDynamicSize
                                                               : d);
  }
}