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