# Copyright 2015 The TensorFlow Authors. All Rights Reserved.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

# =============================================================================

"""Tests for functions."""

from __future__ import absolute_import

from __future__ import division

from __future__ import print_function

import re

import sys

import time

import numpy as np

from tensorflow.core.framework import function_pb2

from tensorflow.core.protobuf import config_pb2

from tensorflow.core.protobuf import rewriter_config_pb2

from tensorflow.python.client import session

from tensorflow.python.framework import constant_op

from tensorflow.python.framework import dtypes

from tensorflow.python.framework import errors_impl

from tensorflow.python.framework import function

from tensorflow.python.framework import graph_to_function_def

from tensorflow.python.framework import ops

from tensorflow.python.framework import tensor_shape

from tensorflow.python.framework import test_util

from tensorflow.python.framework.errors import InvalidArgumentError

from tensorflow.python.ops import array_ops

from tensorflow.python.ops import control_flow_ops

from tensorflow.python.ops import functional_ops

from tensorflow.python.ops import gen_logging_ops

from tensorflow.python.ops import gradients_impl

from tensorflow.python.ops import init_ops

from tensorflow.python.ops import linalg_ops

from tensorflow.python.ops import logging_ops

from tensorflow.python.ops import math_ops

from tensorflow.python.ops import nn_ops

from tensorflow.python.ops import random_ops

from tensorflow.python.ops import variable_scope

from tensorflow.python.ops import variables

from tensorflow.python.platform import test

from tensorflow.python.platform import tf_logging

def _OptimizerOptions():

for cse in [False, True]:

for inline in [False, True]:

for cfold in [False, True]:

cfg = config_pb2.ConfigProto(

graph_options=config_pb2.GraphOptions(

optimizer_options=config_pb2.OptimizerOptions(

opt_level=config_pb2.OptimizerOptions.L0,

do_common_subexpression_elimination=cse,

do_function_inlining=inline,

do_constant_folding=cfold)))

if cse:

cfg.graph_options.rewrite_options.arithmetic_optimization = (

rewriter_config_pb2.RewriterConfig.ON)

else:

cfg.graph_options.rewrite_options.arithmetic_optimization = (

rewriter_config_pb2.RewriterConfig.OFF)

if inline:

cfg.graph_options.rewrite_options.function_optimization = (

rewriter_config_pb2.RewriterConfig.ON)

else:

cfg.graph_options.rewrite_options.function_optimization = (

rewriter_config_pb2.RewriterConfig.OFF)

if cfold:

cfg.graph_options.rewrite_options.constant_folding = (

rewriter_config_pb2.RewriterConfig.ON)

else:

cfg.graph_options.rewrite_options.constant_folding = (

rewriter_config_pb2.RewriterConfig.OFF)

yield cfg

class FunctionTest(test.TestCase):

"""Test methods for verifying Function support.

These test methods are used as mix-ins in two test cases: with

and without C API support.

"""

def testIdentity(self):

@function.Defun(dtypes.float32, func_name="MyIdentity")

def MyIdentityFunc(a):

return a

with ops.Graph().as_default():

call = MyIdentityFunc([18.0])

self.assertEqual("MyIdentity", call.op.name)

with session.Session() as sess:

self.assertAllEqual([18.0], self.evaluate(call))

@test_util.run_v1_only("b/120545219")

def testIdentityImplicitDeref(self):

@function.Defun(dtypes.float32, func_name="MyIdentity")

def MyIdentityFunc(a):

return a

with ops.Graph().as_default():

var = variables.VariableV1([18.0])

call = MyIdentityFunc(var._ref()) # pylint: disable=protected-access

self.assertEqual("MyIdentity", call.op.name)

for cfg in _OptimizerOptions():

with session.Session(config=cfg) as sess:

self.evaluate(var.initializer)

self.assertAllEqual([18.0], self.evaluate(call))

def testIdentityOutputName(self):

@function.Defun(

dtypes.float32, func_name="MyIdentity", out_names=["my_result_name"])

def MyIdentityFunc(a):

return a

with ops.Graph().as_default():

call = MyIdentityFunc([18.0])

self.assertEqual("MyIdentity", call.op.name)

with session.Session() as sess:

self.assertAllEqual([18.0], self.evaluate(call))

def testTooManyOutputNames(self):

@function.Defun(

dtypes.float32,

func_name="MyIdentity",

out_names=["my_result1", "my_result2"])

def MyIdentityFunc(a):

return a

with ops.Graph().as_default():

with self.assertRaisesRegexp(

errors_impl.InvalidArgumentError,

(r"output names must be either empty or equal in size to outputs. "

"output names size = 2 outputs size = 1")):

MyIdentityFunc([18.0])

def testDefineFunction2Args(self):

@function.Defun(dtypes.float32, dtypes.float32, func_name="APlus2B")

def APlus2B(a, b):

return a + b * 2

with ops.Graph().as_default():

call = APlus2B([1.0], [2.0])

self.assertEqual("APlus2B", call.op.name)

with session.Session() as sess:

self.assertAllEqual([5.0], self.evaluate(call))

def testFunctionWithNoOutput(self):

@function.Defun(dtypes.float32, dtypes.float32)

def APlus2B(a, b):

c = a + b * 2 # Create some ops to have nodes in the body

print(c) # Using 'print' to make lint happy

with ops.Graph().as_default():

# Call function. There should be no exceptions.

APlus2B([1.0], [2.0])

def testDefineFunction2ArgsOutputName(self):

@function.Defun(

dtypes.float32,

dtypes.float32,

func_name="APlus2B",

out_names=["my_result_name"])

def APlus2B(a, b):

return a + b * 2

# APlus2B is stateless.

self.assertEqual([], APlus2B.stateful_ops)

with ops.Graph().as_default():

call = APlus2B([1.0], [2.0])

self.assertEqual("APlus2B", call.op.name)

with session.Session() as sess:

self.assertAllEqual([5.0], self.evaluate(call))

def testDefineFunctionDuplicateOutputs(self):

@function.Defun(dtypes.float32, func_name="Duplicate")

def Duplicate(a):

b = a + 1.0

return b, b

g = ops.Graph()

with g.as_default():

Duplicate([3.0])

func_sig = g.as_graph_def().library.function[0].signature

# The names given to both outputs should be different

# even though the same tensor is emitted to both.

out_names = [a.name for a in func_sig.output_arg]

self.assertEqual(2, len(out_names))

self.assertNotEqual(out_names[0], out_names[1])

def testGradientFunc(self):

@function.Defun(dtypes.float32, func_name="XSquarePlusOneFn")

def XSquarePlusOne(x):

return x * x + 1.0

@function.Defun(dtypes.float32, dtypes.float32)

def XSquarePlusOneGrad(x, dy):

dx = functional_ops.symbolic_gradient(

input=[x, dy], Tout=[dtypes.float32], f="XSquarePlusOneFn", name="dx")

return dx

g = ops.Graph()

with g.as_default():

call_f = XSquarePlusOne([2.0])

call_g = XSquarePlusOneGrad([2.0], [0.1])

with session.Session() as sess:

self.assertAllClose([5.0], self.evaluate(call_f))

self.assertAllClose([0.4], self.evaluate(call_g))

def testTanhSymGrad(self):

@function.Defun(dtypes.float32)

def Forward(x):

return math_ops.reduce_sum(math_ops.tanh(x))

g = ops.Graph()

with g.as_default():

x = array_ops.placeholder(dtypes.float32)

y = Forward(x)

dx = gradients_impl.gradients([y], [x])

inp = np.array([-1, 1, 2, -2], dtype=np.float32)

feed = {x: inp}

cfg = config_pb2.ConfigProto(

graph_options=config_pb2.GraphOptions(

optimizer_options=config_pb2.OptimizerOptions(

opt_level=config_pb2.OptimizerOptions.L1,

do_function_inlining=True)))

with session.Session(graph=g, config=cfg) as sess:

out, = sess.run(dx, feed)

self.assertAllClose(1 - np.square(np.tanh(inp)), out)

def testCustomGradient(self):

dtype = dtypes.float32

@function.Defun(dtype, dtype, dtype)

def XentLossGrad(logits, labels, dloss):

dlogits = array_ops.reshape(dloss, [-1, 1]) * (

nn_ops.softmax(logits) - labels)

dlabels = array_ops.zeros_like(labels)

# Takes exp(dlogits) to differentiate it from the "correct" gradient.

return math_ops.exp(dlogits), dlabels

@function.Defun(dtype, dtype, grad_func=XentLossGrad)

def XentLoss(logits, labels):

return math_ops.reduce_sum(labels * math_ops.log(nn_ops.softmax(logits)),

1)

g = ops.Graph()

with g.as_default():

logits = array_ops.placeholder(dtype)

labels = array_ops.placeholder(dtype)

loss = XentLoss(logits, labels)

dlogits = gradients_impl.gradients([loss], [logits])

x = np.random.uniform(-10., 10., size=(4, 9)).astype(np.float32)

prob = np.exp(x) / np.sum(np.exp(x), 1, keepdims=1)

y = np.random.uniform(-10., 10., size=(4, 9)).astype(np.float32)

for cfg in _OptimizerOptions():

tf_logging.info("cfg = %s", cfg)

with session.Session(graph=g, config=cfg) as sess:

out, = sess.run(dlogits, {logits: x, labels: y})

self.assertAllClose(out, np.exp(prob - y))

@test_util.disable_xla("b/124286351") # No error is raised

def testCustomGradientError(self):

dtype = dtypes.float32

@function.Defun(dtype, dtype, dtype)

def Grad(x, dy, dz):

# Should have returned 1 result.

return x, dy + dz

@function.Defun(dtype, grad_func=Grad)

def Forward(x):

return x, x

g = ops.Graph()

with g.as_default():

inp = array_ops.placeholder(dtype)

out = math_ops.add_n(Forward(inp))

dinp = gradients_impl.gradients(out, [inp])

x = np.random.uniform(-10., 10., size=(4, 9)).astype(np.float32)

with session.Session(graph=g) as sess:

with self.assertRaisesRegexp(

errors_impl.InvalidArgumentError,

"SymGrad expects to return 1.*but get 2.*instead"):

_ = sess.run(dinp, {inp: x})

def testSymGradShape(self):

g = ops.Graph()

with g.as_default():

x = array_ops.placeholder(dtypes.float32, [25, 4])

y = array_ops.placeholder(dtypes.float32, [200, 100])

dz = array_ops.placeholder(dtypes.float32, [1])

# We assume Foo is a function of (x, y) -> (z) Then, Foo's

# gradient function is (x, y, dz) -> (dx, dy). dx's shape

# should be the same as x's; and dy's shape should be the same

# as y's.

dx, dy = functional_ops.symbolic_gradient(

input=[x, y, dz], Tout=[dtypes.float32] * 2, f="Foo")

self.assertEqual(x.get_shape(), dx.get_shape())

self.assertEqual(y.get_shape(), dy.get_shape())

@test_util.run_deprecated_v1

def testSymGradAttr(self):

@function.Defun(noinline=True)

def Foo(x):

return x * 2

self.assertTrue(

Foo.instantiate([dtypes.float32]).definition.attr["_noinline"].b)

g = ops.Graph()

with g.as_default():

x = constant_op.constant(3.0)

y = Foo(x)

dx, = gradients_impl.gradients(y, [x])

cfg = config_pb2.ConfigProto(

graph_options=config_pb2.GraphOptions(

optimizer_options=config_pb2.OptimizerOptions(

opt_level=config_pb2.OptimizerOptions.L0,

do_common_subexpression_elimination=True,

do_function_inlining=True,

do_constant_folding=True)))

with self.session(graph=g, config=cfg):

self.assertAllClose(y.eval(), 6.)

self.assertAllClose(dx.eval(), 2.)

def _testZNoDepOnY(self, use_const_grad_ys):

@function.Defun(dtypes.float32, dtypes.float32)

def Foo(x, y): # pylint: disable=unused-argument

return x * 2

with ops.Graph().as_default():

# z = Foo(x, y). z doe

x = constant_op.constant(1.0)

y = constant_op.constant(2.0)

z = Foo(x, y)

if use_const_grad_ys:

dx, dy = gradients_impl.gradients([z], [x, y], grad_ys=[1.0])

else:

dx, dy = gradients_impl.gradients([z], [x, y])

with session.Session() as sess:

dx_val, dy_val = self.evaluate([dx, dy])

self.assertEqual([2.0], dx_val)

self.assertEqual([0.0], dy_val)

def testZNoDepOnY(self):

self._testZNoDepOnY(False)

def testZNoDepOnYConstGradYs(self):

# Tests for constant folding of grad_ys

self._testZNoDepOnY(True)

def testDefineFunctionNoArgs(self):

@function.Defun(func_name="AConstant")

def AConstant():

return constant_op.constant([42])

with ops.Graph().as_default():

call = AConstant()

self.assertEqual("AConstant", call.op.name)

with session.Session() as sess:

self.assertAllEqual([42], self.evaluate(call))

def testDefineFunctionNames(self):

@function.Defun(dtypes.float32, func_name="Foo")

def Foo(a):

return a + 1

with ops.Graph().as_default():

call1 = Foo([1.0])

self.assertEqual("Foo", call1.op.name)

call2 = Foo([1.0])

self.assertEqual("Foo_1", call2.op.name)

# pylint: disable=unexpected-keyword-arg

call3 = Foo([1.0], name="mine")

self.assertEqual("mine", call3.op.name)

with ops.name_scope("my"):

call4 = Foo([1.0], name="precious")

self.assertEqual("my/precious", call4.op.name)

def testNoOp(self):

@function.Defun(dtypes.float32)

def Foo(x):

y = logging_ops.Print(x, [], "Hello")

with ops.control_dependencies([y]):

z = control_flow_ops.no_op()

with ops.control_dependencies([z]):

return x * 2

with ops.Graph().as_default(), self.cached_session():

z = Foo(constant_op.constant(3.0))

self.assertAllEqual(z.eval(), 6.0)

def testAssertOp(self):

@function.Defun(dtypes.float32)

def Foo(x):

check = gen_logging_ops._assert(math_ops.greater(x, 0), [x])

with ops.control_dependencies([check]):

return x * 2

# Foo contains a stateful op (Assert).

self.assertEqual([("Assert", "Assert")], Foo.stateful_ops)

g = ops.Graph()

with g.as_default(), self.cached_session():

self.assertAllEqual(Foo(constant_op.constant(3.0)).eval(), 6.0)

with self.assertRaisesRegexp(errors_impl.InvalidArgumentError,

"assertion failed.*-3"):

self.assertAllEqual(Foo(constant_op.constant(-3.0)).eval(), 6.0)

@test_util.run_deprecated_v1

def testAssertWrapper(self):

@function.Defun(dtypes.float32)

def MyFn(x):

with ops.control_dependencies(

[control_flow_ops.Assert(math_ops.less_equal(x, 10.0), [x])]):

return array_ops.identity(x)

with self.cached_session():

self.assertEqual(1.0, MyFn(1.0).eval())

with self.assertRaisesRegexp(errors_impl.InvalidArgumentError,

"assertion"):

_ = MyFn(100.0).eval()

@test_util.run_deprecated_v1

def testWhileLoopCallsFunc(self):

with self.session(use_gpu=True) as sess:

@function.Defun(dtypes.float32)

def Times2(x):

constant_two = constant_op.constant(2, dtypes.int32)

two_on_gpu = math_ops.cast(constant_two, dtypes.float32)

return x * two_on_gpu

def Body(x):

x2 = Times2(x)

x2.set_shape([])

return x2

loop = control_flow_ops.while_loop(lambda x: x < 1e5, Body, [1.0])

ans = self.evaluate(loop)

self.assertAllClose(ans, 131072.)

@test_util.run_deprecated_v1

def testControlFlowStrictness(self):

"""Inlined functions must not execute in a untaken control flow branch."""

@function.Defun(dtypes.int32)

def AssertFail(x):

# Assertion that always fails and does not have a data dependency on `x`.

assert_false = control_flow_ops.Assert(False, [42])

with ops.control_dependencies([assert_false]):

return array_ops.identity(x)

with ops.device("CPU"):

pred = array_ops.placeholder(dtypes.bool)

x = array_ops.placeholder(dtypes.int32)

cond = control_flow_ops.cond(pred, lambda: x + 1, lambda: AssertFail(x))

# pylint: disable=unnecessary-lambda

loop = control_flow_ops.while_loop(lambda y: pred,

lambda y: AssertFail(y), [x])

# pylint: enable=unnecessary-lambda

rewriter_config = rewriter_config_pb2.RewriterConfig(

dependency_optimization=rewriter_config_pb2.RewriterConfig.OFF)

# Enables inlining.

config = config_pb2.ConfigProto(

graph_options=config_pb2.GraphOptions(

optimizer_options=config_pb2.OptimizerOptions(

opt_level=config_pb2.OptimizerOptions.L0,

do_common_subexpression_elimination=True,

do_function_inlining=True,

do_constant_folding=True),

rewrite_options=rewriter_config))

with session.Session(config=config) as sess:

# Since the 'False' branch is not taken, the assertion should not fire.

self.assertEqual(4, sess.run(cond, {pred: True, x: 3}))

# The assertion should still fire if the False branch is taken.

with self.assertRaisesRegexp(errors_impl.InvalidArgumentError,

"assertion"):

sess.run(cond, {pred: False, x: 3})

# Similarly for loops.

self.assertEqual(3, sess.run(loop, {pred: False, x: 3}))

with self.assertRaisesRegexp(errors_impl.InvalidArgumentError,

"assertion"):

sess.run(loop, {pred: True, x: 3})

@test_util.run_deprecated_v1

def testVar(self):

@function.Defun(dtypes.float32)

def Foo(x):

return x * x + 1

g = ops.Graph()

with g.as_default():

v = variables.Variable(constant_op.constant(10.0))

z = Foo(v)

with self.session(graph=g):

variables.global_variables_initializer().run()

self.assertAllEqual(z.eval(), 101.)

@test_util.run_deprecated_v1

def testResourceVarAsImplicitInput(self):

g = ops.Graph()

with g.as_default(), ops.device("cpu:0"):

expected_type = dtypes.float32

expected_shape = tensor_shape.TensorShape((4, 4))

v = variable_scope.get_variable(

"var", expected_shape, expected_type, use_resource=True)

@function.Defun()

def Foo():

captured = array_ops.identity(v)

self.assertEqual(expected_type, captured.dtype)

self.assertEqual(expected_shape, captured.shape)

return captured, array_ops.shape(captured)

expected_val = v.value()

actual_val, actual_shape = Foo()

with self.session(graph=g):

v.initializer.run()

self.assertAllEqual(expected_val.eval(), self.evaluate(actual_val))

self.assertAllEqual(expected_shape, self.evaluate(actual_shape))

def testDefineErrors(self):

with ops.Graph().as_default():

with self.assertRaisesRegexp(ValueError, "can not return None"):

@function.Defun()

def TwoNone():

return None, None

_ = TwoNone.definition

with self.assertRaisesRegexp(ValueError, "are not supported"):

@function.Defun()

def DefaultArg(unused_a=12):

return constant_op.constant([1])

_ = DefaultArg.definition

with self.assertRaisesRegexp(ValueError, "are not supported"):

@function.Defun()

def KwArgs(**unused_kwargs):

return constant_op.constant([1])

_ = KwArgs.definition

with self.assertRaisesRegexp(ValueError, "specified input types"):

@function.Defun(dtypes.float32)

def PlusMinusV2(a, b):

return a + b, b - a

_ = PlusMinusV2.definition

with self.assertRaisesRegexp(ValueError, "specified input types"):

@function.Defun(dtypes.float32, dtypes.float32, dtypes.float32)

def PlusMinusV3(a, b):

return a + b, b - a

_ = PlusMinusV3.definition

def testCallErrors(self):

@function.Defun()

def Const():

return constant_op.constant(1)

@function.Defun(dtypes.int32)

def PlusOne(a):

return a + 1

@function.Defun(dtypes.int32, dtypes.int32)

def PlusMinus(a, b):

return a + b, b - a

with ops.Graph().as_default():

_ = Const()

# pylint: disable=too-many-function-args

# pylint: disable=unexpected-keyword-arg

# pylint: disable=no-value-for-parameter

with self.assertRaisesRegexp(ValueError, "arguments: 0"):

_ = Const(1)

with self.assertRaisesRegexp(ValueError, "arguments: 0"):

_ = Const(1, 2)

with self.assertRaisesRegexp(ValueError, "arguments: 1"):

_ = PlusOne()

_ = PlusOne(1)

with self.assertRaisesRegexp(ValueError, "arguments: 1"):

_ = PlusOne(1, 2)

with self.assertRaisesRegexp(ValueError, "arguments: 2"):

_ = PlusMinus()

with self.assertRaisesRegexp(ValueError, "arguments: 2"):

_ = PlusMinus(1)

_ = PlusMinus(1, 2)

_ = PlusOne(1, name="p1")

with self.assertRaisesRegexp(ValueError, "Unknown keyword arguments"):

_ = PlusOne(1, device="/device:GPU:0")

def testFunctionDecorator(self):

@function.Defun(dtypes.float32, func_name="Minus1")

def Minus1(b):

return b - 1.0

with ops.Graph().as_default():

call1 = Minus1([2.])

self.assertTrue(isinstance(Minus1, function._DefinedFunction))

self.assertEqual(Minus1.name, "Minus1")

# pylint: disable=unexpected-keyword-arg

call2 = Minus1(call1, name="next")

# pylint: enable=unexpected-keyword-arg

self.assertEqual("next", call2.op.name)

with session.Session() as sess:

self.assertAllEqual([1], self.evaluate(call1))

self.assertAllEqual([0], self.evaluate(call2))

def testNestedFunction(self):

@function.Defun(dtypes.float32)

def Cube(x):

return x * x * x

@function.Defun(dtypes.float32, dtypes.float32)

def CubeXPlusY(x, y):

return Cube(x) + y

with ops.Graph().as_default():

z = CubeXPlusY(3.0, -2.0)

with self.cached_session():

self.assertAllEqual(z.eval(), 25.0)

def testNestedDefinedFunction(self):

@function.Defun(dtypes.float32, dtypes.float32)

def CubeXPlusY(x, y):

@function.Defun(dtypes.float32)

def Cube(x):

return x * x * x

return Cube(x) + y

with ops.Graph().as_default():

z = CubeXPlusY(3.0, -2.0)

with self.cached_session():

self.assertAllEqual(z.eval(), 25.0)

def testUnusedFunction(self):

invoked = False

# pylint: disable=unused-variable

@function.Defun()

def Unused():

invoked = True

return constant_op.constant(42.)

self.assertFalse(invoked)

g = ops.Graph()

with g.as_default():

@function.Defun()

def Unused2():

invoked = True

return constant_op.constant(7.)

constant_op.constant(3.)

# pylint: enable=unused-variable

self.assertFalse(invoked)

gdef = g.as_graph_def()

self.assertEqual(0, len(gdef.library.function))

@test_util.run_deprecated_v1

def testReduction(self):

g = ops.Graph()

# BN0 is computing batch normed matrix along rows.

def BN0(x):

mean = math_ops.reduce_mean(x, [0])

var = math_ops.reduce_mean(math_ops.square(x - mean)) # biased var

rstd = math_ops.rsqrt(var + 1e-8)

return (x - mean) * rstd

# Wraps BatchNorm in a tf function.

@function.Defun(dtypes.float32)

def BN1(x):

return BN0(x)

with g.as_default():

x = array_ops.placeholder(dtypes.float32)

y0 = BN0(x) # A plain graph

y1 = BN1(x) # A tf function

dx0, = gradients_impl.gradients([y0], [x])

dx1, = gradients_impl.gradients([y1], [x])

# Both should produce the same result and gradient.

with self.session(graph=g) as sess:

vals = sess.run([y0, y1, dx0, dx1], {x: np.random.uniform(size=(3, 7))})

self.assertAllClose(vals[0], vals[1])

self.assertAllClose(vals[2], vals[3])

@test_util.run_deprecated_v1

def testCapture(self):

g = ops.Graph()

with g.as_default():

w = variables.Variable(constant_op.constant([[1.0]]))

b = variables.Variable(constant_op.constant([2.0]))

# Foo() captures w and b.

@function.Defun(dtypes.float32)

def Foo(x):

# Plus() captures b.

@function.Defun(dtypes.float32)

def Plus(y):

return y + b

return Plus(math_ops.matmul(w, x))

y = Foo(constant_op.constant([[10.]]))

@function.Defun()

def Bar():

return w

z = Bar()

with self.session(graph=g):

variables.global_variables_initializer().run()

self.assertAllEqual(y.eval(), [[12.0]])

self.assertAllEqual(z.eval(), [[1.0]])

def testCaptureControls(self):

g = ops.Graph()

with g.as_default():

x = constant_op.constant([10.0])

x = logging_ops.Print(x, [x], "outer")

@function.Defun(dtypes.float32)

def Foo(y):

with ops.control_dependencies([x]):

y = logging_ops.Print(y, [y], "inner")

return y

with self.assertRaisesRegexp(ValueError, "not an element of this graph."):

# NOTE: We still do not support capturing control deps.

_ = Foo(x)

@test_util.run_deprecated_v1

def testCaptureInWhileLoop(self):

g = ops.Graph()

with g.as_default():

x = constant_op.constant(1)

@function.Defun()

def Foo():

return control_flow_ops.while_loop(lambda i: i < 10, lambda i: i + x,

[0])

y = Foo()

with self.session(graph=g) as sess:

self.assertEqual(self.evaluate(y), 10)

@test_util.run_deprecated_v1

def testCaptureInCond(self):

g = ops.Graph()

with g.as_default():

x = constant_op.constant(1)

@function.Defun(dtypes.bool)

def Foo(pred):

return control_flow_ops.cond(pred, lambda: x, lambda: x + 1)

y = Foo(True)

z = Foo(False)

with self.session(graph=g) as sess:

self.assertEqual(self.evaluate(y), 1)

self.assertEqual(self.evaluate(z), 2)

def testStableName(self):

@function.Defun()

def Foo(x, y, z):

return math_ops.tanh(math_ops.matmul(x, y) + z)

if sys.byteorder == "big":

self.assertEqual("Foo_kEdkAG8SJvg",

Foo.instantiate([dtypes.float32] * 3).name)

else:

self.assertEqual("Foo_aCYSbwBkR5A",

Foo.instantiate([dtypes.float32] * 3).name)

@test_util.run_deprecated_v1

def testSignatureHash(self):

# Foo.Inner and Bar.Inner have identical function body but have

# different signatures. They should be treated as two different functions.

@function.Defun()

def Foo(x):

@function.Defun()

def Inner(x):

return x + 10.

return Inner(x)

@function.Defun()

def Bar(x):

@function.Defun()

def Inner(x, unused_y, unused_z):

return x + 10.

return Inner(x, 2., 3.)

g = ops.Graph()

with g.as_default():

x = constant_op.constant(10.0)

y = Foo(x)

z = Bar(x)

with self.session(graph=g) as sess:

v0, v1 = self.evaluate([y, z])

self.assertAllEqual(v0, 20.)

self.assertAllEqual(v1, 20.)

def testShapeFunction(self):

@function.Defun(

dtypes.float32, shape_func=lambda op: [op.inputs[0].get_shape()])

def Foo(x):

return x + 1.0

@function.Defun(

shape_func=lambda op: [[1] + op.inputs[0].get_shape().as_list()])

def Bar(x):

return array_ops.stack([x])

g = ops.Graph()

with g.as_default():

x = Foo([1.0, 2.0])

self.assertEqual(x.get_shape().as_list(), [2])

y = Bar(array_ops.zeros([1, 2, 3]))

self.assertAllEqual(y.get_shape().as_list(), [1, 1, 2, 3])

@test_util.run_deprecated_v1

def testVariableReuse(self):

def LinearWithReuse(input_tensor, reuse=None):

size = input_tensor.shape.dims[1]

with variable_scope.variable_scope("linear", reuse=reuse):

w = variable_scope.get_variable(

"w", shape=[size, size], dtype=input_tensor.dtype)

return math_ops.matmul(input_tensor, w)

@function.Defun(dtypes.float32)

def Foo(inputs):

inputs = array_ops.reshape(inputs, [32, 100])

hidden = LinearWithReuse(inputs)

return LinearWithReuse(hidden, reuse=True)

input_op = array_ops.placeholder(shape=[32, 100], dtype=dtypes.float32)

output_op = Foo(input_op)

global_vars = variables.global_variables()

self.assertEqual(len(global_vars), 1)

self.assertEqual(global_vars[0].name, "linear/w:0")

with session.Session() as sess:

self.evaluate(variables.global_variables_initializer())

output_val = sess.run(

output_op, feed_dict={input_op: np.random.rand(32, 100)})

self.assertEqual(output_val.shape, (32, 100))

@test_util.run_deprecated_v1

def testFunctionCallInDifferentVariableScopes(self):

@function.Defun(dtypes.float32)

def Foo(inputs):

var = variable_scope.get_variable(

"var",

shape=[10],

dtype=dtypes.float32,

initializer=init_ops.ones_initializer())

return inputs + var

input_op = array_ops.placeholder(shape=[10], dtype=dtypes.float32)

with variable_scope.variable_scope("vs1"):

out1_op = Foo(input_op)

with variable_scope.variable_scope("vs2"):

out2_op = Foo(input_op)

global_vars = variables.global_variables()

self.assertEqual(len(global_vars), 1)

self.assertEqual(global_vars[0].name, "vs1/var:0")

with session.Session() as sess:

self.evaluate(variables.global_variables_initializer())

out1, out2 = sess.run(

[out1_op, out2_op], feed_dict={input_op: np.linspace(1, 10, 10)})

self.assertAllEqual(out1, np.linspace(2, 11, 10))

self.assertAllEqual(out2, np.linspace(2, 11, 10))

def testTwoInputsSameOp(self):

g = ops.Graph()

with g.as_default():

m = array_ops.placeholder(dtypes.float32)

s, u, v = linalg_ops.svd(m)

ss = math_ops.reduce_sum(s)

uu = math_ops.reduce_sum(u)

vv = math_ops.reduce_sum(v)

result = ss + uu + vv

f = graph_to_function_def.graph_to_function_def(

g,

g.get_operations()[1:], # skip the placeholder

[s, u, v],

[result])

self.assertEqual(len(f.signature.input_arg), 3)

def testGradientWithIntegerFunctionArgument(self):

@function.Defun(dtypes.int32, dtypes.float32)

def Foo(t, x):

return x[t]

g = ops.Graph()

with g.as_default():

inp = array_ops.placeholder(dtypes.float32)

t = constant_op.constant(0, dtypes.int32)

out = Foo(t, inp)

dinp, = gradients_impl.gradients(out, [inp])

x = np.zeros((2,)).astype(np.float32)

with session.Session(graph=g) as sess:

self.assertAllClose(

np.array([1.0, 0.0]).astype(np.float32), sess.run(dinp, {inp: x}))

@test_util.run_deprecated_v1

def testFunctionMarkedStateful(self):

@function.Defun(dtypes.int32, dtypes.float32)

def Foo(t, x):

return x[t]

@function.Defun(dtypes.int64)

def Bar(x):

return x

# NOTE(mrry): All functions are currently considered stateless by the

# runtime, so we simulate a "stateful" function.

# TODO(b/70565970): Remove this hack when we are able to build stateful