list->add_s(nested_value.s());
          break;
        case AttrValue::kType:
          list->add_type(nested_value.type());
          break;
        case AttrValue::kShape:
          *list->add_shape() = nested_value.shape();
          break;
        default:
          return errors::Unimplemented("Unhandled nested attribute!");
      }
    } else if (auto attr = mlir::dyn_cast<mlir::ElementsAttr>(a)) {

                       AttrValueMap* values) {
  AttrValue value;
  auto& shape_list = *value.mutable_list();
  for (const std::optional<llvm::ArrayRef<int64_t>>& shape : shapes) {
    TensorShapeProto& tshape = *shape_list.add_shape();
    if (shape.has_value()) {
      for (int64_t dim : *shape) {
        tshape.add_dim()->set_size(mlir::ShapedType::isDynamic(dim) ? -1 : dim);
      }
    } else {
      tshape.set_unknown_rank(true);
    }

        handle_dtypes_attr.mutable_list()->add_type(dtype);

        SetTensorShapeProto(subtype,
                            handle_shapes_attr.mutable_list()->add_shape());
      }

      (*node_def->mutable_attr())["_handle_dtypes"] = handle_dtypes_attr;
      (*node_def->mutable_attr())["_handle_shapes"] = handle_shapes_attr;
    }
  }

  TF_RETURN_IF_ERROR(

        handle_dtypes_attr.mutable_list()->add_type(dtype);

        SetTensorShapeProto(subtype,
                            handle_shapes_attr.mutable_list()->add_shape());
      }

      (*node_def->mutable_attr())["_handle_dtypes"] = handle_dtypes_attr;
      (*node_def->mutable_attr())["_handle_shapes"] = handle_shapes_attr;
    }
  }

  TF_RETURN_IF_ERROR(

        const InferredShape& arg_shape = it->second;
        shape_inference::InferenceContext* context =
            shape_refiner->GetContext(n);

        if (arg_shape.handle_type != DT_INVALID) {
          shape_inference::ShapeHandle handle;
          TF_RETURN_IF_ERROR(context->MakeShapeFromPartialTensorShape(
              arg_shape.handle_shape, &handle));

};

struct VariantSubTypeBuilder {
  typedef VariantSubType Table;
  ::flatbuffers::FlatBufferBuilder &fbb_;
  ::flatbuffers::uoffset_t start_;
  void add_shape(::flatbuffers::Offset<::flatbuffers::Vector<int32_t>> shape) {
    fbb_.AddOffset(VariantSubType::VT_SHAPE, shape);
  }
  void add_type(tflite::TensorType type) {

absl::StatusOr<TensorShape> GetTensorShapeFromXlaArgument(
    const XlaArgument& arg) {
  if (absl::holds_alternative<xla::Shape>(arg.shape)) {
    TensorShape arg_shape;
    TF_RETURN_IF_ERROR(
        XLAShapeToTensorShape(std::get<xla::Shape>(arg.shape), &arg_shape));
    return arg_shape;
  } else {
    return std::get<TensorShape>(arg.shape);
  }
}

Status MaybeRewriteLayoutWithShardedShape(
    mlir::StringAttr sharding,

  // TODO(kramm): Move this into a separate pass. See b/184944903
  xla::Shape old_shape = xla::TypeToShape(t);
  XLA_Shape old_shape_c = {};
  XLA_Shape new_shape_c = {};
  TfTpu_ExecutorApiFn *executor = stream_executor::tpu::ExecutorApiFn();
  if (!stream_executor::tpu::IsInitialized(executor)) {
    return failure();
  }
  ApiConverter::ToC(old_shape, &old_shape_c);

  auto paddings_value = builder->create<TF::ConstOp>(filter.getLoc(), paddings);
  std::vector<int64_t> pad_shape = {filter_shape[0] + pad_h,
                                    filter_shape[1] + pad_w, filter_shape[2],
                                    filter_shape[3]};
  SmallVector<int64_t, 4> expand_shape(pad_shape.begin(), pad_shape.end());

  auto expand_result_type =

      return failure();
    }
    // It is already a 1-D constant, no change.
    auto old_shape = shape.getShapedType().getShape();
    if (old_shape.size() == 1) {
      return failure();
    }
    // Verify all the leading dimensions are length one, except the last one.
    for (auto it = ++old_shape.rbegin(); it != old_shape.rend(); ++it) {
      if (*it != 1) {
        reshape->emitError(