Method Compile

tensorflow/compiler/tf2xla/kernels/arg_op.cc:37–66

  }

  void Compile(XlaOpKernelContext* ctx) override {
    // If 'frame' is non-null, this is a function call inside an outer JIT
    // compilation. Use the usual implementation of _Arg.
    auto frame = ctx->call_frame();
    if (frame != nullptr) {
      Tensor val;
      OP_REQUIRES_OK(ctx, frame->GetArg(index_, &val));
      // Types that cannot be copied using memcpy (like DT_STRING) are wrapped
      // in a DT_UINT8 and hence the type mismatches. Skip the test in such
      // cases. See XlaOpKernelContext::SetOutputExpression for details.
      if (DataTypeCanUseMemcpy(dtype_)) {
        OP_REQUIRES(ctx, val.dtype() == dtype_,
                    errors::InvalidArgument(
                        "Type mismatch: actual ", DataTypeString(val.dtype()),
                        " vs. expect ", DataTypeString(dtype_)));
      }
      // Forwards the argument from the frame.
      ctx->op_kernel_context()->set_output(0, val);
      return;
    }

    const XlaExpression& arg = ctx->xla_context()->args()[index_];
    OP_REQUIRES(ctx, arg.kind() != XlaExpression::Kind::kInvalid,
                errors::InvalidArgument("Invalid/missing argument expression"));
    if (ctx->expected_output_dtype(0) == DT_VARIANT) {
      ctx->SetTensorListOutput(0, arg.handle());
    } else {
      ctx->SetOutputExpression(0, arg);
    }
  }

 private:
  int index_;