TODO(b/35949885): There is duplication here with the handling of the while_op. Refactor the common code out/rework.
| 72 | // TODO(b/35949885): There is duplication here with the handling of the |
| 73 | // while_op. Refactor the common code out/rework. |
| 74 | void XlaIfOp::Compile(XlaOpKernelContext* ctx) { |
| 75 | xla::XlaBuilder* b = ctx->builder(); |
| 76 | |
| 77 | OP_REQUIRES(ctx, cond_type_ == DT_BOOL, |
| 78 | errors::InvalidArgument( |
| 79 | "Condition argument must be a boolean for XLA compilation")); |
| 80 | OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(ctx->InputShape(0)), |
| 81 | errors::InvalidArgument( |
| 82 | "Condition argument must be a scalar for XLA compilation")); |
| 83 | |
| 84 | VLOG(1) << "Building If: " << input_types_.size() << " inputs"; |
| 85 | |
| 86 | std::vector<XlaCompiler::Argument> arguments(input_types_.size()); |
| 87 | int num_resource_args = 0; |
| 88 | for (int i = 0; i < input_types_.size(); ++i) { |
| 89 | XlaCompiler::Argument& arg = arguments[i]; |
| 90 | DataType type = ctx->input_type(i + 1); |
| 91 | |
| 92 | if (type == DT_RESOURCE) { |
| 93 | XlaResource* resource; |
| 94 | OP_REQUIRES_OK(ctx, ctx->GetResourceInput(i + 1, &resource)); |
| 95 | |
| 96 | arg.initialized = resource->initialized(); |
| 97 | arg.kind = XlaCompiler::Argument::kResource; |
| 98 | arg.resource_kind = resource->kind(); |
| 99 | |
| 100 | arg.type = resource->type(); |
| 101 | arg.shape = resource->shape(); |
| 102 | OP_REQUIRES(ctx, arg.initialized, |
| 103 | errors::Unimplemented("Uninitialized arguments: ", arg.name)); |
| 104 | arg.max_array_size = resource->max_array_size(); |
| 105 | for (const auto& gradient : resource->tensor_array_gradients()) { |
| 106 | arg.tensor_array_gradients.insert(gradient.first); |
| 107 | } |
| 108 | arg.name = resource->name(); |
| 109 | VLOG(2) << "Resource " << resource->name() |
| 110 | << " type: " << DataTypeString(arg.type) |
| 111 | << " shape: " << arg.HumanString() |
| 112 | << " initialized: " << arg.initialized; |
| 113 | |
| 114 | num_resource_args++; |
| 115 | } else { |
| 116 | arg.kind = XlaCompiler::Argument::kParameter; |
| 117 | arg.type = input_types_[i]; |
| 118 | // Use the xla::Shape for the input instead of ctx->InputShape. This is |
| 119 | // necessary for forwarding shapes of DT_VARIANTs, e.g. TensorLists. |
| 120 | auto shape_or = ctx->builder()->GetShape(ctx->Input(i + 1)); |
| 121 | OP_REQUIRES_OK(ctx, shape_or.status()); |
| 122 | arg.shape = shape_or.ValueOrDie(); |
| 123 | VLOG(2) << "Arg type: " << DataTypeString(arg.type) |
| 124 | << " shape: " << arg.HumanString(); |
| 125 | } |
| 126 | } |
| 127 | |
| 128 | if (propagate_compile_time_consts_) { |
| 129 | // Replaces `kParameter` type args in `arguments` with `kConstant` if |
| 130 | // the op input corresponding to that arg is a compile-time const. This |
| 131 | // is necessary to propagate compile time consts to ops in the branch |
no test coverage detected