| 339 | } |
| 340 | |
| 341 | void RenderInterfaceImpl(const OpSpec& op, RenderMode mode, |
| 342 | SourceWriter* writer) { |
| 343 | ArgumentSpec output = op.outputs().front(); |
| 344 | |
| 345 | if (mode == OPERAND) { |
| 346 | bool cast2obj = output.type().wildcard(); |
| 347 | Type return_type = |
| 348 | Type::Class("Output", "org.tensorflow") |
| 349 | .add_parameter(cast2obj ? Type::Class("Object") : output.type()); |
| 350 | Method as_output = Method::Create("asOutput", return_type) |
| 351 | .add_annotation(Annotation::Create("Override")); |
| 352 | if (cast2obj) { |
| 353 | as_output.add_annotation( |
| 354 | Annotation::Create("SuppressWarnings").attributes("\"unchecked\"")); |
| 355 | } |
| 356 | writer->BeginMethod(as_output, PUBLIC); |
| 357 | if (cast2obj) { |
| 358 | writer->Append("return (").AppendType(return_type).Append(") "); |
| 359 | } else { |
| 360 | writer->Append("return "); |
| 361 | } |
| 362 | writer->Append(output.var().name() + ";").EndLine().EndMethod(); |
| 363 | |
| 364 | } else if (mode == LIST_OPERAND) { |
| 365 | Type operand = Type::Interface("Operand", "org.tensorflow"); |
| 366 | if (output.type().wildcard()) { |
| 367 | operand.add_parameter(Type::Class("Object")); |
| 368 | } else { |
| 369 | operand.add_parameter(output.type()); |
| 370 | } |
| 371 | Type return_type = |
| 372 | Type::Interface("Iterator", "java.util").add_parameter(operand); |
| 373 | Method iterator = |
| 374 | Method::Create("iterator", return_type) |
| 375 | .add_annotation(Annotation::Create("Override")) |
| 376 | .add_annotation(Annotation::Create("SuppressWarnings") |
| 377 | .attributes("{\"rawtypes\", \"unchecked\"}")); |
| 378 | // cast the output list using a raw List |
| 379 | writer->BeginMethod(iterator, PUBLIC) |
| 380 | .Append("return (" + return_type.name() + ") ") |
| 381 | .Append(output.var().name() + ".iterator();") |
| 382 | .EndLine() |
| 383 | .EndMethod(); |
| 384 | } |
| 385 | } |
| 386 | |
| 387 | void RenderOptionsClass(const OpSpec& op, const Type& op_class, |
| 388 | SourceWriter* writer) { |