From 949b18f0fedc5ccdd48b37e6308402a688212fb1 Mon Sep 17 00:00:00 2001
From: Chris Jones <cjfj@google.com>
Date: Wed, 18 Dec 2019 07:16:34 -0800
Subject: [PATCH] Fix reduction with a specified axis.

PiperOrigin-RevId: 286185937
Change-Id: I56bed77a84b758f204888102e8a7d7175b6e3d9e
---
 tensorflow/python/distribute/distribute_lib.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/distribute/distribute_lib.py b/tensorflow/python/distribute/distribute_lib.py
index e988499292e..216ec8b62d2 100644
--- a/tensorflow/python/distribute/distribute_lib.py
+++ b/tensorflow/python/distribute/distribute_lib.py
@@ -840,7 +840,7 @@ class Strategy(object):
         if dim is not None:
           # By returning a python value in the static shape case, we can
           # maybe get a fast path for reducing the denominator.
-          return numer, dim
+          return numer, array_ops.constant(dim, dtype=dtypes.int64)
       elif axis < 0:
         axis = axis + array_ops.rank(v)
       if v.shape.rank == 1: