| 26 | namespace tensorflow { |
| 27 | |
| 28 | XlaIfOp::XlaIfOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { |
| 29 | const NameAttrList* name_attr; |
| 30 | OP_REQUIRES_OK(ctx, ctx->GetAttr("then_branch", &name_attr)); |
| 31 | then_branch_ = *name_attr; |
| 32 | OP_REQUIRES_OK(ctx, ctx->GetAttr("else_branch", &name_attr)); |
| 33 | else_branch_ = *name_attr; |
| 34 | |
| 35 | OP_REQUIRES_OK(ctx, ctx->GetAttr("Tcond", &cond_type_)); |
| 36 | OP_REQUIRES_OK(ctx, ctx->GetAttr("Tin", &input_types_)); |
| 37 | OP_REQUIRES_OK(ctx, ctx->GetAttr("Tout", &output_types_)); |
| 38 | if (!ctx->GetAttr(kXlaTokenInputNodesAttrName, &token_input_nodes_).ok()) { |
| 39 | has_token_input_output_ = false; |
| 40 | } else { |
| 41 | has_token_input_output_ = !token_input_nodes_.empty(); |
| 42 | } |
| 43 | if (ctx->HasAttr(kPropagateCompileTimeConsts)) { |
| 44 | OP_REQUIRES_OK(ctx, ctx->GetAttr(kPropagateCompileTimeConsts, |
| 45 | &propagate_compile_time_consts_)); |
| 46 | } |
| 47 | } |
| 48 | |
| 49 | Status ConvertCompileTimeConstArgumentsToConst( |
| 50 | XlaOpKernelContext* ctx, std::vector<XlaCompiler::Argument>* args) { |