for cse in [False, True]:

for inline in [False, True]:

for cfold in [False, True]:

yield 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,

"""Test methods for verifying Function 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], sess.run(call))

def testIdentityOutputName(self):

@function.Defun(

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], sess.run(call))

def testTooManyOutputNames(self):

@function.Defun(

def MyIdentityFunc(a):

return a

with ops.Graph().as_default():

with self.assertRaisesRegexp(

ValueError, (r"Length of out_names \(2\) does not match number of "

r"outputs \(1\): my_result1, my_result2")):

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], sess.run(call))

def testValueErrorOnFunctionWithNoOutput(self):

# TODO(iga): Remove this restriction and this test

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

def APlus2B(a, b):

print(a + b * 2) # Create some ops to have nodes in the body

# Using 'print' to make lint happy

with ops.Graph().as_default():

with self.assertRaisesRegexp(ValueError,

"Function can not return None"):

APlus2B([1.0], [2.0])

def testDefineFunction2ArgsOutputName(self):

@function.Defun(

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], sess.run(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], sess.run(call_f))

self.assertAllClose([0.4], sess.run(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))

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())

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.test_session(graph=g, config=cfg):

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

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

def testZNoDepOnY(self):

@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)

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

with session.Session() as sess:

dx_val, dy_val = sess.run([dx, dy])

self.assertEqual([2.0], dx_val)

self.assertEqual([0.0], dy_val)

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], sess.run(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, [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.test_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

g = ops.Graph()

with g.as_default(), self.test_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.disable_c_api # Op._add_control_inputs doesn't work with C API

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.test_session():

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

with self.assertRaisesRegexp(errors_impl.InvalidArgumentError,

"assertion"):

_ = MyFn(100.0).eval()

@test_util.disable_c_api # Op._add_control_inputs doesn't work with C API

def testWhileLoopCallsFunc(self):

with self.test_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 = sess.run(loop)

self.assertAllClose(ans, 131072.)

@test_util.disable_c_api # Op._add_control_inputs doesn't work with C API

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

# 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)))

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})

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.test_session(graph=g):

variables.global_variables_initializer().run()

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

def testResourceVarAsImplicitInput(self):

g = ops.Graph()

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

v = variable_scope.get_variable(

"var", (4, 4), dtypes.float32, use_resource=True)

@function.Defun()

def Foo():

return array_ops.identity(v)

y = v.value()

z = Foo()

with self.test_session(graph=g):

v.initializer.run()

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

def testDefineErrors(self):

with ops.Graph().as_default():

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

@function.Defun()

def NoResult():

pass

_ = NoResult.definition

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], sess.run(call1))

self.assertAllEqual([0], sess.run(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.test_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.test_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))

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.test_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])

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.]]))

with self.test_session(graph=g):

variables.global_variables_initializer().run()

self.assertAllEqual(y.eval(), [[12.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)

def testStableName(self):

@function.Defun()

def Foo(x, y, z):

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

self.assertEqual("Foo_d643acf7", Foo.instantiate([dtypes.float32] * 3).name)

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.test_session(graph=g) as sess:

v0, v1 = sess.run([y, z])

self.assertAllEqual(v0, 20.)

self.assertAllEqual(v1, 20.)

def testShapeFunction(self):

@function.Defun(

def Foo(x):

return x + 1.0

@function.Defun(

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])

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:

sess.run(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))

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:

sess.run(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])

def expectFunctionsEqual(self, func, grad_func=None, new_func=None):

if new_func is None:

# Make a copy of func.definition to avoid any bugs masked by using the

# same object

serialized_fdef = func.definition.SerializeToString()

# Serialize and then deserialize `func` to create `new_func`

fdef = function_pb2.FunctionDef.FromString(serialized_fdef)

new_func = function._from_definition(fdef, grad_func=grad_func)

self.assertEqual(func.name, new_func.name)

self.assertEqual(func.definition, new_func.definition)

self.assertEqual(func.grad_func_name, new_func.grad_func_name)

self.assertEqual(func.declared_input_types, new_func.declared_input_types)

self.assertEqual(func.captured_inputs, new_func.captured_inputs)

def testBasic(self):

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

def Foo(x, y):

return x + y

self.expectFunctionsEqual(Foo)

def testGradFunc(self):

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

def G(x, dy):

return x * dy

@function.Defun(dtypes.float32, grad_func=G)

def F(x):

return math_ops.exp(x) - math_ops.exp(-x)

self.expectFunctionsEqual(F, grad_func=G)

def testCapturedInputs(self):

c = constant_op.constant(10, dtypes.int64)

@function.Defun(dtypes.int64)

def Foo(x):

return x + c

new_func = function._from_definition(Foo.definition)

self.assertEqual(Foo.name, new_func.name)

self.assertEqual(Foo.definition, new_func.definition)

self.assertEqual(Foo.grad_func_name, new_func.grad_func_name)

# Captured inputs are added as regular inputs to the function definition

self.assertEqual(new_func.declared_input_types,

Foo.declared_input_types + (dtypes.int64,))

self.assertEqual(len(new_func.captured_inputs), 0)

def testNestedFunctions(self):

@function.Defun(dtypes.float32)

def Outer(x):

@function.Defun(dtypes.float32)

def Inner(y):

return y + 1

return Inner(Inner(x))

self.expectFunctionsEqual(Outer)

def testFromLibrary(self):

# Define some functions with different gradient functions. Note that many of

# the below functions are identical since function bodies don't matter for

# this test.

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

def G1(x, dy):

return x * dy

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

def G2(x, dy):

return x * dy

# F1 and F2 have the same gradient function

@function.Defun(dtypes.float32, grad_func=G1)

def F1(x):

return math_ops.exp(x) - math_ops.exp(-x)

@function.Defun(dtypes.float32, grad_func=G1)

def F2(x):

return math_ops.exp(x) - math_ops.exp(-x)

# F3 has a different gradient function

@function.Defun(dtypes.float32, grad_func=G2)

def F3(x):

return math_ops.exp(x) - math_ops.exp(-x)

# F4 has no gradient function

@function.Defun(dtypes.float32)

def F4(x):

return math_ops.exp(x) - math_ops.exp(-x)

# Instantiate all functions

g = ops.Graph()

with g.as_default():

c = constant_op.constant(1.0, dtypes.float32)

f1 = F1(c)

f2 = F2(c)

f3 = F3(c)

f4 = F4(c)

gradients_impl.gradients([f1, f2, f3, f4], c)

library = g.as_graph_def().library

new_funcs = function._from_library(library)

def CheckNewFunc(func):

new_func = [f for f in new_funcs if f.name == func.name]

self.assertEqual(len(new_func), 1)

self.expectFunctionsEqual(func, new_func=new_func[0])

CheckNewFunc(G1)

CheckNewFunc(G2)

CheckNewFunc(F1)

CheckNewFunc(F2)

CheckNewFunc(F3)

CheckNewFunc(F4)