From 742778f04e0d28495025fc27051d861de67b285a Mon Sep 17 00:00:00 2001
From: Yuefeng Zhou <yuefengz@google.com>
Date: Wed, 26 Dec 2018 12:17:17 -0800
Subject: [PATCH] Enable collective ops in defun

PiperOrigin-RevId: 226942243
---
 tensorflow/core/kernels/partitioned_function_ops.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/core/kernels/partitioned_function_ops.cc b/tensorflow/core/kernels/partitioned_function_ops.cc
index fbecd909bea..96ade7900f3 100644
--- a/tensorflow/core/kernels/partitioned_function_ops.cc
+++ b/tensorflow/core/kernels/partitioned_function_ops.cc
@@ -439,6 +439,7 @@ class PartitionedCallOp : public AsyncOpKernel {
     opts.step_container = ctx->step_container();
     opts.cancellation_manager = ctx->cancellation_manager();
     opts.stats_collector = ctx->stats_collector();
+    opts.collective_executor = ctx->collective_executor();
     // TODO(akshayka): Consider selecting a runner on a per-device basis, i.e.,
     // using device-specific threadpools when available.
     opts.runner = ctx->runner();