| 111 | }; |
| 112 | |
| 113 | Status TensorProducerImpl::Init( |
| 114 | OpKernelContext* ctx, const NameAttrList& func) { |
| 115 | func_ = func; |
| 116 | env_ = ctx->env(); |
| 117 | runner_ = *ctx->runner(); |
| 118 | TF_RETURN_IF_ERROR(ctx->function_library()->Clone(&flib_def_, &pflr_, &lib_)); |
| 119 | FunctionLibraryRuntime::InstantiateOptions inst_opts; |
| 120 | inst_opts.state_handle = std::to_string(random::New64()); |
| 121 | TF_RETURN_IF_ERROR( |
| 122 | lib_->Instantiate(func_.name(), AttrSlice(&func_.attr()), |
| 123 | inst_opts, &f_handle_)); |
| 124 | const FunctionBody* fbody = lib_->GetFunctionBody(f_handle_); |
| 125 | if (fbody == nullptr) { |
| 126 | return errors::Internal("Failed to instantiate function body."); |
| 127 | } |
| 128 | if (fbody->ret_types.size() != 1) { |
| 129 | return errors::InvalidArgument( |
| 130 | "Tensor Generator Func should return 1 tensor"); |
| 131 | } |
| 132 | return Status::OK(); |
| 133 | } |
| 134 | |
| 135 | void TensorProducerImpl::Produce(const TensorGenerator::Consumer& consumer) { |
| 136 | TensorProducerContext* ctx = new TensorProducerContext; |