"Only tensors and tuples of tensors are supported right now.");
  }
}

size_t GetFlatSize(const TaggedValue& value) {
  if (value.type() == TaggedValue::Type::TUPLE) {
    size_t result = 0;
    for (const auto& t : value.tuple()) {
      result += GetFlatSize(t);
    }
    return result;
  } else if (value.type() == TaggedValue::Type::LIST) {
    size_t result = 0;