def testIdentityShape(self): with self.cached_session(): shape = [2, 3] array_2x3 = [[1, 2, 3], [6, 5, 4]] tensor = constant_op.constant(array_2x3) self.assertEquals(shape, tensor.get_shape()) self.assertEquals(shape, array_ops.identity_n([tensor])[0].get_shape()) self.assertEquals(shape, array_ops.identity_n([array_2x3])[0].get_shape())
def testIdentityShape(self): with self.test_session(): shape = [2, 3] array_2x3 = [[1, 2, 3], [6, 5, 4]] tensor = constant_op.constant(array_2x3) self.assertEquals(shape, tensor.get_shape()) self.assertEquals(shape, array_ops.identity_n([tensor])[0].get_shape()) self.assertEquals(shape, array_ops.identity_n([array_2x3])[0].get_shape())
def prune(self, feeds, fetches): flat_feeds, flat_fetches = nest.flatten(feeds), nest.flatten(fetches) for f in flat_feeds + flat_fetches: if not isinstance(f, ops.Tensor): raise ValueError("Feeds and fetches must be tensors.") if f.graph is not self._func_graph: raise ValueError( "Can only prune function whose feeds and fetches " "are from this graph (%s). Tensor %s from graph %s" % (self._func_graph, f, f.graph)) with self._func_graph.as_default(): pruned_graph = func_graph.FuncGraph("pruned") sink_tensor = array_ops.identity_n(flat_fetches)[0] lift_map = lift_to_graph.lift_to_graph(sink_tensor, pruned_graph, sources=flat_feeds + self.graph.internal_captures) pruned_graph.outputs.extend(lift_map[x] for x in flat_fetches) for external_capture, internal_capture in self.graph.captures.items(): pruned_graph.captures[external_capture] = lift_map[ internal_capture] pruned_graph.inputs.extend(lift_map[x] for x in flat_feeds) pruned_graph.inputs.extend(pruned_graph.captures.values()) pruned_graph.structured_outputs = nest.map_structure( lambda node: lift_map[node], fetches) pruned_fn = WrappedFunction(pruned_graph, variable_holder=self._variable_holder) pruned_fn._num_positional_args = len(flat_feeds) # pylint: disable=protected-access pruned_fn._arg_keywords = [] # pylint: disable=protected-access return pruned_fn
def testCopiesOfUnsupportedTypesFailGracefully(self): """Tests that copies of unsupported types don't crash.""" test_types = set([ np.uint8, np.uint16, np.uint32, np.uint64, np.int8, np.int16, np.int32, np.int64, np.float16, np.float32, np.float16, dtypes.bfloat16.as_numpy_dtype ]) shape = (10, 10) for unsupported_dtype in test_types - self.all_types: with self.test_session() as sess: with ops.device("CPU"): x = array_ops.placeholder(unsupported_dtype, shape) with self.test_scope(): y, = array_ops.identity_n([x]) with ops.device("CPU"): z = array_ops.identity(y) inputs = np.random.randint(-100, 100, shape) inputs = inputs.astype(unsupported_dtype) # Execution should either succeed or raise an InvalidArgumentError, # but not crash. Even "unsupported types" may succeed here since some # backends (e.g., the CPU backend) are happy to handle buffers of # unsupported types, even if they cannot compute with them. try: sess.run(z, {x: inputs}) except errors.InvalidArgumentError: pass
def outer(y, shp): y, shp = control_flow_ops.while_loop_v2( lambda *_: True, inner, (y, shp), maximum_iterations=3) y, shp = array_ops.identity_n([y, shp]) return control_flow_ops.while_loop_v2(lambda *_: True, inner, (y, shp), maximum_iterations=5)
def testCopiesOfUnsupportedTypesFailGracefully(self): """Tests that copies of unsupported types don't crash.""" test_types = set([ np.uint8, np.uint16, np.uint32, np.uint64, np.int8, np.int16, np.int32, np.int64, np.float16, np.float32, np.float16, dtypes.bfloat16.as_numpy_dtype ]) shape = (10, 10) for unsupported_dtype in test_types - self.all_types: with self.session() as sess: with ops.device("CPU"): x = array_ops.placeholder(unsupported_dtype, shape) with self.test_scope(): y, = array_ops.identity_n([x]) with ops.device("CPU"): z = array_ops.identity(y) inputs = np.random.randint(-100, 100, shape) inputs = inputs.astype(unsupported_dtype) # Execution should either succeed or raise an InvalidArgumentError, # but not crash. Even "unsupported types" may succeed here since some # backends (e.g., the CPU backend) are happy to handle buffers of # unsupported types, even if they cannot compute with them. try: sess.run(z, {x: inputs}) except errors.InvalidArgumentError: pass
def Run(branch, x, fetch_by_name, use_gpu=use_gpu): with ops.Graph().as_default() as g: @function.Defun(dtypes.float32) def two(x): return -1, x * 2 @function.Defun(dtypes.float32) def three(x): return 0, x * 3 @function.Defun(dtypes.float32) def four(x): return 1, x * 4 outputs = gen_functional_ops.case(branch, input=[x], Tout=[dtypes.int32, dtypes.float32], branches=[two, three, four], name="my_case") # `outputs` is the list of output tensors of the Case op. We # arbitrarily choose the 0th tensor to get the Case op and set the # lowering attribute on it. outputs[0].op._set_attr("_lower_using_switch_merge", attr_value_pb2.AttrValue(b=True)) outputs = array_ops.identity_n(outputs) with self.session(graph=g, use_gpu=use_gpu) as sess: return sess.run("my_case:1" if fetch_by_name else outputs[1])
def testInt32String_6(self): with self.test_session() as sess: [value0, value1] = sess.run( array_ops.identity_n([[1, 2, 3, 4, 5, 6], [b"a", b"b", b"C", b"d", b"E", b"f", b"g"]])) self.assertAllEqual(np.array([1, 2, 3, 4, 5, 6]), value0) self.assertAllEqual( np.array([b"a", b"b", b"C", b"d", b"E", b"f", b"g"]), value1)
def testInt32String_6(self): with self.cached_session() as sess: [value0, value1] = sess.run( array_ops.identity_n([[1, 2, 3, 4, 5, 6], [b"a", b"b", b"C", b"d", b"E", b"f", b"g"]])) self.assertAllEqual(np.array([1, 2, 3, 4, 5, 6]), value0) self.assertAllEqual( np.array([b"a", b"b", b"C", b"d", b"E", b"f", b"g"]), value1)
def testInt32String_6(self): value0, value1 = self.evaluate( array_ops.identity_n([[1, 2, 3, 4, 5, 6], [b"a", b"b", b"C", b"d", b"E", b"f", b"g"]])) self.assertAllEqual(np.array([1, 2, 3, 4, 5, 6]), value0) self.assertAllEqual( np.array([b"a", b"b", b"C", b"d", b"E", b"f", b"g"]), value1)
def prune(self, feeds, fetches): flat_feeds, flat_fetches = nest.flatten(feeds), nest.flatten(fetches) for f in flat_feeds: if not isinstance(f, ops.Tensor): raise ValueError("Feeds must be tensors.") tensor_fetches = [] operation_fetches = [] for f in flat_fetches: if isinstance(f, ops.Tensor): tensor_fetches.append(f) elif isinstance(f, ops.Operation): operation_fetches.append(f) else: raise ValueError("Fetches must be tensors or operations.") for f in flat_feeds + flat_fetches: if f.graph is not self._func_graph: raise ValueError( "Can only prune function whose feeds and fetches " "are from this graph (%s). Tensor %s from graph %s" % (self._func_graph, f, f.graph)) with self._func_graph.as_default(): pruned_graph = func_graph.FuncGraph("pruned") with ops.control_dependencies(operation_fetches): if tensor_fetches: identity_fetches = array_ops.identity_n(tensor_fetches) sink_tensor = identity_fetches[0] else: identity_fetches = [] sink_tensor = control_flow_ops.no_op() lift_map = lift_to_graph.lift_to_graph(sink_tensor, pruned_graph, sources=flat_feeds + self.graph.internal_captures) for original_fetch, identity_fetch in zip(tensor_fetches, identity_fetches): lift_map[original_fetch] = lift_map[identity_fetch] pruned_graph.outputs.extend(lift_map[x] for x in flat_fetches if isinstance(x, ops.Tensor)) for external_capture, internal_capture in self.graph.captures.items(): pruned_graph.captures[external_capture] = lift_map[ internal_capture] pruned_graph.inputs.extend(lift_map[x] for x in flat_feeds) pruned_graph.inputs.extend(pruned_graph.captures.values()) def _structured_output_mapping(fetched): lifted = lift_map[fetched] if isinstance(lifted, ops.Operation): return None return lifted pruned_graph.structured_outputs = nest.map_structure( _structured_output_mapping, fetches) pruned_fn = WrappedFunction(pruned_graph, variable_holder=self._variable_holder) pruned_fn._num_positional_args = len(flat_feeds) # pylint: disable=protected-access pruned_fn._arg_keywords = [] # pylint: disable=protected-access return pruned_fn
def testFastpathExecute_IdentityNCorrectResponse(self): ctx = context.context() a_2_by_2 = random_ops.random_uniform((2, 2)) b_2_by_2 = random_ops.random_uniform((2, 2)) self.assertAllClose( array_ops.identity_n([a_2_by_2, b_2_by_2]), pywrap_tensorflow.TFE_Py_FastPathExecute(ctx._handle, ctx.device_name, "IdentityN", None, None, [a_2_by_2, b_2_by_2]))
def testFastpathExecute_IdentityNCorrectResponse(self): ctx = context.context() a_2_by_2 = random_ops.random_uniform((2, 2)) b_2_by_2 = random_ops.random_uniform((2, 2)) self.assertAllClose( array_ops.identity_n([a_2_by_2, b_2_by_2]), pywrap_tensorflow.TFE_Py_FastPathExecute(ctx._handle, ctx.device_name, "IdentityN", None, None, [a_2_by_2, b_2_by_2]))
def fwd_gradients(ys, xs, grad_xs=None, assert_unused=False): """Computes forward-mode derivatives. This is accomplished in pure-python using tensorflow's existing (reverse-mode) gradients. There is additional overhead on graph construction, but runtime performance should be equal to a manual implementation [citation needed]. See https://j-towns.github.io/2017/06/12/A-new-trick.html and https://github.com/HIPS/autograd/pull/175 for the original discussion of this method, and https://github.com/renmengye/tensorflow-forward-ad for a "direct" implementation. Args: ys: A list of tensors. xs: A list of tensors. grad_xs: An optional list of tensors. If provided, must have the same length and shapes compatible with xs. assert_unused: Add assertions that intermediate values are not computed. Returns: A list of tensors of the same shapes as ys. The directional derivatives of ys with respect to xs in the direction grad_xs. Leaving grad_xs unspecified is equivalent to passing in 1s for each x in xs. """ # This version of forward-mode autodiff is based on code by Tim Cooijmans # and handles list arguments and certain special cases such as when the # ys doesn't depend on one or more of the xs, and when tf.IndexedSlices are # generated by the first tf.gradients call. us = [array_ops.zeros_like(y) + float('nan') for y in ys] dydxs = gradients(ys, xs, grad_ys=us) # deal with strange types that tf.gradients returns but can't deal with dydxs = [ ops.convert_to_tensor(dydx) if isinstance(dydx, ops.IndexedSlices) else dydx for dydx in dydxs ] if assert_unused: with ops.control_dependencies(dydxs): assert_unused = control_flow_ops.Assert(False, [1], name='fwd_gradients') with ops.control_dependencies([assert_unused]): dydxs = array_ops.identity_n(dydxs) dydxs = [ array_ops.zeros_like(x) if dydx is None else dydx for x, dydx in zip(xs, dydxs) ] for x, dydx in zip(xs, dydxs): dydx.set_shape(x.shape) dysdx = gradients(dydxs, us, grad_ys=grad_xs) return dysdx
def testFastpathExecute_IdentityNCorrectResponse(self): ctx = context.context() ctx.ensure_initialized() a_2_by_2 = random_ops.random_uniform((2, 2)) b_2_by_2 = random_ops.random_uniform((2, 2)) self.assertAllClose( array_ops.identity_n([a_2_by_2, b_2_by_2]), pywrap_tfe.TFE_Py_FastPathExecute(ctx, "IdentityN", None, [a_2_by_2, b_2_by_2]))
def testInt32_shapes(self): with self.cached_session() as sess: inp0 = constant_op.constant([10, 20, 30, 40, 50, 60], shape=[2, 3]) inp1 = constant_op.constant([11, 21, 31, 41, 51, 61], shape=[3, 2]) inp2 = constant_op.constant( [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15], shape=[5, 3]) [value0, value1, value2] = sess.run(array_ops.identity_n([inp0, inp1, inp2])) self.assertAllEqual(np.array([[10, 20, 30], [40, 50, 60]]), value0) self.assertAllEqual(np.array([[11, 21], [31, 41], [51, 61]]), value1) self.assertAllEqual( np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12], [13, 14, 15]]), value2)
def testInt32_shapes(self): inp0 = constant_op.constant([10, 20, 30, 40, 50, 60], shape=[2, 3]) inp1 = constant_op.constant([11, 21, 31, 41, 51, 61], shape=[3, 2]) inp2 = constant_op.constant( [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15], shape=[5, 3]) value0, value1, value2 = self.evaluate( array_ops.identity_n([inp0, inp1, inp2])) self.assertAllEqual(np.array([[10, 20, 30], [40, 50, 60]]), value0) self.assertAllEqual(np.array([[11, 21], [31, 41], [51, 61]]), value1) self.assertAllEqual( np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12], [13, 14, 15]]), value2)
def fwd_gradients(ys, xs, grad_xs=None, assert_unused=False): """Computes forward-mode derivatives. This is accomplished in pure-python using tensorflow's existing (reverse-mode) gradients. There is additional overhead on graph construction, but runtime performance should be equal to a manual implementation [citation needed]. See https://j-towns.github.io/2017/06/12/A-new-trick.html and https://github.com/HIPS/autograd/pull/175 for the original discussion of this method, and https://github.com/renmengye/tensorflow-forward-ad for a "direct" implementation. Args: ys: A list of tensors. xs: A list of tensors. grad_xs: An optional list of tensors. If provided, must have the same length and shapes compatible with xs. assert_unused: Add assertions that intermediate values are not computed. Returns: A list of tensors of the same shapes as ys. The directional derivatives of ys with respect to xs in the direction grad_xs. Leaving grad_xs unspecified is equivalent to passing in 1s for each x in xs. """ # This version of forward-mode autodiff is based on code by Tim Cooijmans # and handles list arguments and certain special cases such as when the # ys doesn't depend on one or more of the xs, and when tf.IndexedSlices are # generated by the first tf.gradients call. us = [array_ops.zeros_like(y) + float('nan') for y in ys] dydxs = gradients(ys, xs, grad_ys=us) # deal with strange types that tf.gradients returns but can't deal with dydxs = [ops.convert_to_tensor(dydx) if isinstance(dydx, ops.IndexedSlices) else dydx for dydx in dydxs] if assert_unused: with ops.control_dependencies(dydxs): assert_unused = control_flow_ops.Assert(False, [1], name='fwd_gradients') with ops.control_dependencies([assert_unused]): dydxs = array_ops.identity_n(dydxs) dydxs = [array_ops.zeros_like(x) if dydx is None else dydx for x, dydx in zip(xs, dydxs)] for x, dydx in zip(xs, dydxs): dydx.set_shape(x.shape) dysdx = gradients(dydxs, us, grad_ys=grad_xs) return dysdx
def decorated(*args, **kwargs): """Decorated function with custom gradient.""" if context.in_graph_mode(): if kwargs: raise ValueError( "custom_gradient in graph mode doesn't support keyword arguments." ) name = "CustomGradient-%s" % tf_ops.uid() args = [tf_ops.convert_to_tensor(x) for x in args] result, grad_fn = f(*args) flat_result = nest.flatten(result) all_tensors = flat_result + args @tf_ops.RegisterGradient(name) def internal_grad_fn(unused_op, *result_grads): # pylint: disable=unused-variable gradients = nest.flatten( grad_fn(*result_grads[:len(flat_result)])) # Need to return one value per input to the IdentityN, so pad the # gradients of the inputs of the custom_gradient function with the # gradients of the outputs as well. return ([None] * len(flat_result)) + gradients with tf_ops.get_default_graph().gradient_override_map( {"IdentityN": name}): all_tensors = array_ops.identity_n(all_tensors) return nest.pack_sequence_as( structure=result, flat_sequence=all_tensors[:len(flat_result)]) input_tensors = [] for x in args: if isinstance(x, tf_ops.Tensor): input_tensors.append(x) if isinstance(x, resource_variable_ops.ResourceVariable): input_tensors.append(x.read_value()) with tape.stop_recording(): result, grad_fn = f(*args, **kwargs) # TODO(apassos): naive uses of custom_gradient will not get the correct # second derivative this way if they capture any output tensors. Change the # signature of custom_gradient. def actual_grad_fn(*outputs): return nest.flatten(grad_fn(*outputs)) flat_result = nest.flatten(result) tape.record_operation(f.__name__, flat_result, input_tensors, actual_grad_fn) flat_result = list(flat_result) return result
def testFastpathExecute_IdentityNTapeWrite(self): ctx = context.context() a_2_by_2 = random_ops.random_uniform((2, 2)) b_2_by_2 = random_ops.random_uniform((2, 2)) with backprop.GradientTape(persistent=True) as tape: tape.watch(a_2_by_2) tape.watch(b_2_by_2) z1 = pywrap_tensorflow.TFE_Py_FastPathExecute( ctx._handle, ctx.device_name, "IdentityN", None, None, [a_2_by_2, b_2_by_2]) z2 = array_ops.identity_n([a_2_by_2, b_2_by_2]) dz1_dy = tape.gradient(z1[0], [a_2_by_2])[0] dz2_dy = tape.gradient(z2[0], [a_2_by_2])[0] self.assertAllEqual(dz1_dy.numpy(), dz2_dy.numpy())
def testFastpathExecute_IdentityNTapeWrite(self): ctx = context.context() a_2_by_2 = random_ops.random_uniform((2, 2)) b_2_by_2 = random_ops.random_uniform((2, 2)) with backprop.GradientTape(persistent=True) as tape: tape.watch(a_2_by_2) tape.watch(b_2_by_2) z1 = pywrap_tensorflow.TFE_Py_FastPathExecute( ctx._handle, ctx.device_name, "IdentityN", None, None, [a_2_by_2, b_2_by_2]) z2 = array_ops.identity_n([a_2_by_2, b_2_by_2]) dz1_dy = tape.gradient(z1[0], [a_2_by_2])[0] dz2_dy = tape.gradient(z2[0], [a_2_by_2])[0] self.assertAllEqual(dz1_dy.numpy(), dz2_dy.numpy())
def decorated(*args, **kwargs): """Decorated function with custom gradient.""" if context.in_graph_mode(): if kwargs: raise ValueError( "custom_gradient in graph mode doesn't support keyword arguments.") name = "CustomGradient-%s" % tf_ops.uid() args = [tf_ops.convert_to_tensor(x) for x in args] result, grad_fn = f(*args) flat_result = nest.flatten(result) all_tensors = flat_result + args @tf_ops.RegisterGradient(name) def internal_grad_fn(unused_op, *result_grads): # pylint: disable=unused-variable gradients = nest.flatten(grad_fn(*result_grads[:len(flat_result)])) # Need to return one value per input to the IdentityN, so pad the # gradients of the inputs of the custom_gradient function with the # gradients of the outputs as well. return ([None] * len(flat_result)) + gradients with tf_ops.get_default_graph().gradient_override_map( {"IdentityN": name}): all_tensors = array_ops.identity_n(all_tensors) return nest.pack_sequence_as( structure=result, flat_sequence=all_tensors[:len(flat_result)]) input_tensors = [x for x in args if isinstance(x, tf_ops.Tensor)] with tape.stop_recording(): result, grad_fn = f(*args, **kwargs) # TODO(apassos): naive uses of custom_gradient will not get the correct # second derivative this way if they capture any output tensors. Change the # signature of custom_gradient. def actual_grad_fn(*outputs): return grad_fn(*outputs) flat_result = nest.flatten(result) tape.record_operation( f.__name__, flat_result, input_tensors, [], actual_grad_fn) flat_result = list(flat_result) return result
def decorated(*args, **kwargs): """Decorated function with custom gradient.""" if context.in_graph_mode(): if kwargs: raise ValueError( "custom_gradient in graph mode doesn't support keyword arguments.") name = "CustomGradient-%s" % tf_ops.uid() args = [tf_ops.convert_to_tensor(x) for x in args] result, grad_fn = f(*args) flat_result = nest.flatten(result) all_tensors = flat_result + args @tf_ops.RegisterGradient(name) def internal_grad_fn(unused_op, *result_grads): # pylint: disable=unused-variable gradients = nest.flatten(grad_fn(*result_grads[:len(flat_result)])) # Need to return one value per input to the IdentityN, so pad the # gradients of the inputs of the custom_gradient function with the # gradients of the outputs as well. return ([None] * len(flat_result)) + gradients with tf_ops.get_default_graph().gradient_override_map( {"IdentityN": name}): all_tensors = array_ops.identity_n(all_tensors) return nest.pack_sequence_as( structure=result, flat_sequence=all_tensors[:len(flat_result)]) input_tensors = [tf_ops.convert_to_tensor(x) for x in args] result, grad_fn = f(*args, **kwargs) flat_result = nest.flatten(result) # TODO(apassos) consider removing the identity below. flat_result = [gen_array_ops.identity(x) for x in flat_result] def actual_grad_fn(*outputs): return nest.flatten(grad_fn(*outputs)) tape.record_operation( f.__name__, flat_result, input_tensors, actual_grad_fn) flat_result = list(flat_result) return nest.pack_sequence_as(result, flat_result)
def decorated(*args, **kwargs): """Decorated function with custom gradient.""" if context.in_graph_mode(): if kwargs: raise ValueError( "custom_gradient in graph mode doesn't support keyword arguments." ) name = "CustomGradient-%s" % tf_ops.uid() args = [tf_ops.convert_to_tensor(x) for x in args] result, grad_fn = f(*args) flat_result = nest.flatten(result) all_tensors = flat_result + args @tf_ops.RegisterGradient(name) def internal_grad_fn(unused_op, *result_grads): # pylint: disable=unused-variable gradients = nest.flatten( grad_fn(*result_grads[:len(flat_result)])) # Need to return one value per input to the IdentityN, so pad the # gradients of the inputs of the custom_gradient function with the # gradients of the outputs as well. return ([None] * len(flat_result)) + gradients with tf_ops.get_default_graph().gradient_override_map( {"IdentityN": name}): all_tensors = array_ops.identity_n(all_tensors) return nest.pack_sequence_as( structure=result, flat_sequence=all_tensors[:len(flat_result)]) input_tensors = [tf_ops.convert_to_tensor(x) for x in args] with tape.stop_recording(): result, grad_fn = f(*args, **kwargs) flat_result = nest.flatten(result) # TODO(apassos) consider removing the identity below. flat_result = [gen_array_ops.identity(x) for x in flat_result] def actual_grad_fn(*outputs): return nest.flatten(grad_fn(*outputs)) tape.record_operation(f.__name__, flat_result, input_tensors, actual_grad_fn) flat_result = list(flat_result) return nest.pack_sequence_as(result, flat_result)
def fwd_gradients(ys, xs, grad_xs=None, assert_unused=False): us = [array_ops.zeros_like(y) + float('nan') for y in ys] dydxs = tf.gradients(ys, xs, grad_ys=us) dydxs = [ops.convert_to_tensor(dydx) if isinstance(dydx, ops.IndexedSlices) else dydx for dydx in dydxs] if assert_unused: with ops.control_dependencies(dydxs): assert_unused = control_flow_ops.Assert(False, [1], name='fwd_gradients') with ops.control_dependencies([assert_unused]): dydxs = array_ops.identity_n(dydxs) dydxs = [array_ops.zeros_like(x) if dydx is None else dydx for x, dydx in zip(xs, dydxs)] for x, dydx in zip(xs, dydxs): dydx.set_shape(x.shape) dysdx = tf.gradients(dydxs, us, grad_ys=grad_xs) return dysdx
def prune(self, feeds, fetches): flat_feeds, flat_fetches = nest.flatten(feeds), nest.flatten(fetches) for f in flat_feeds + flat_fetches: if not isinstance(f, ops.Tensor): raise ValueError("Feeds and fetches must be tensors.") if f.graph is not self._func_graph: raise ValueError( "Can only prune function whose feeds and fetches " "are from this graph (%s). Tensor %s from graph %s" % ( self._func_graph, f, f.graph)) with self._func_graph.as_default(): pruned_graph = func_graph.FuncGraph("pruned") sink_tensor = array_ops.identity_n(flat_fetches)[0] lift_map = lift_to_graph.lift_to_graph( sink_tensor, pruned_graph, sources=flat_feeds) pruned_graph.outputs.extend(lift_map[x] for x in flat_fetches) pruned_graph.inputs.extend(lift_map[x] for x in flat_feeds) pruned_fn = WrappedFunction( pruned_graph, variable_holder=self._variable_holder) pruned_fn._num_positional_args = len(flat_feeds) # pylint: disable=protected-access pruned_fn._arg_keywords = [] # pylint: disable=protected-access return pruned_fn
def Test(): x = tf.constant(1.0, shape=(5, 3)) y = tf.constant(1.0, shape=(3, 5)) s = tf.matmul(x, y) t = tf.matmul(y, x) [t, s] = array_ops.identity_n([t, s]) tensor_info_x = tf.compat.v1.saved_model.utils.build_tensor_info(x) tensor_info_y = tf.compat.v1.saved_model.utils.build_tensor_info(y) tensor_info_s = tf.compat.v1.saved_model.utils.build_tensor_info(s) tensor_info_t = tf.compat.v1.saved_model.utils.build_tensor_info(t) return { 'key': (tf.compat.v1.saved_model.signature_def_utils.build_signature_def( inputs={ 'x': tensor_info_x, 'y': tensor_info_y }, outputs={ 's': tensor_info_s, 't': tensor_info_t }, method_name='some_function')), 'key2': (tf.compat.v1.saved_model.signature_def_utils.build_signature_def( inputs={ 'a': tensor_info_y, 'b': tensor_info_x, }, outputs={ 'c': tensor_info_t, 'd': tensor_info_s, }, method_name='reverse_arguments')) }
def Run(branch, x): @function.Defun(dtypes.float32) def two(x): return -1, x * 2 @function.Defun(dtypes.float32) def three(x): return 0, x * 3 @function.Defun(dtypes.float32) def four(x): return 1, x * 4 outputs = gen_functional_ops.case(branch, input=[x], Tout=[dtypes.int32, dtypes.float32], branches=[two, three, four]) # `outputs` is the list of output tensors of the Case op. We # arbitrarily choose the 0th tensor to get the Case op and set the # lowering attribute on it. outputs[0].op._set_attr("_lower_using_switch_merge", attr_value_pb2.AttrValue(b=True)) outputs = array_ops.identity_n(outputs) return outputs[1]
def _graph_mode_decorator(f, *args, **kwargs): """Implement custom gradient decorator for graph mode.""" # TODO(rsepassi): Add support for kwargs if kwargs: raise ValueError( "The custom_gradient decorator currently supports keywords " "arguments only when eager execution is enabled.") name = "CustomGradient-%s" % ops.uid() args = [ops.convert_to_tensor(x) for x in args] # Checking global and local variables attempts to ensure that no non-resource # Variables are added to the graph. current_var_scope = variable_scope.get_variable_scope() before_vars = set(current_var_scope.global_variables() + current_var_scope.local_variables()) with backprop.GradientTape() as tape: result, grad_fn = f(*args) after_vars = set(current_var_scope.global_variables() + current_var_scope.local_variables()) new_vars = after_vars - before_vars for v in new_vars: if not isinstance(v, resource_variable_ops.ResourceVariable): raise TypeError( "All variables used by a function wrapped with @custom_gradient must " "be `ResourceVariable`s. Ensure that no `variable_scope` is created " "with `use_resource=False`.") # The variables that grad_fn needs to return gradients for are the set of # variables used that are *not* part of the inputs. variables = list(set(tape.watched_variables()) - set(args)) grad_argspec = tf_inspect.getfullargspec(grad_fn) variables_in_signature = ("variables" in grad_argspec.args or grad_argspec.varkw) if variables and not variables_in_signature: raise TypeError("If using @custom_gradient with a function that " "uses variables, then grad_fn must accept a keyword " "argument 'variables'.") if variables_in_signature and not variables: # User seems to intend to use variables but none were captured. if not variable_scope.get_variable_scope().use_resource: raise TypeError( "If using @custom_gradient with a function that " "uses variables, the enclosing variable scope must " "have use_resource=True.") else: logging.warn( "@custom_gradient grad_fn has 'variables' in signature, but " "no ResourceVariables were used on the forward pass.") flat_result = nest.flatten(result) all_tensors = flat_result + args + variables @ops.RegisterGradient(name) def internal_grad_fn(unused_op, *result_grads): # pylint: disable=unused-variable """Custom grad fn wrapper.""" result_grads = result_grads[:len(flat_result)] if variables: input_grads, variable_grads = grad_fn(*result_grads, variables=variables) if len(variable_grads) != len(variables): raise ValueError("Must return gradient for each variable from " "@custom_gradient grad_fn.") else: input_grads = grad_fn(*result_grads) variable_grads = [] # Need to return one value per input to the IdentityN, so pad the # gradients of the inputs of the custom_gradient function with the # gradients of the outputs as well. input_grads = nest.flatten(input_grads) return ([None] * len(flat_result)) + input_grads + variable_grads original_tensors = all_tensors with ops.get_default_graph().gradient_override_map({"IdentityN": name}): all_tensors = array_ops.identity_n(all_tensors) for ot, t in zip(original_tensors, all_tensors): copy_handle_data(ot, t) return nest.pack_sequence_as(structure=result, flat_sequence=all_tensors[:len(flat_result)])
def _graph_mode_decorator(f, args, kwargs): """Implement custom gradient decorator for graph mode.""" # TODO(rsepassi): Add support for kwargs if kwargs: raise ValueError( "The custom_gradient decorator currently supports keywords " "arguments only when eager execution is enabled.") name = "CustomGradient-%s" % ops.uid() args = nest.map_structure(ops.convert_to_tensor, args) # Checking global and local variables attempts to ensure that no non-resource # Variables are added to the graph. current_var_scope = variable_scope.get_variable_scope() before_vars = set([ v.ref() for v in current_var_scope.global_variables() + current_var_scope.local_variables() ]) with tape_lib.VariableWatcher() as variable_watcher: result, grad_fn = f(*args) args = nest.flatten(args) after_vars = set([ v.ref() for v in current_var_scope.global_variables() + current_var_scope.local_variables() ]) new_vars = after_vars - before_vars new_vars_list = [v.deref() for v in new_vars] for v in new_vars_list: if not resource_variable_ops.is_resource_variable(v): raise TypeError( "All variables used by a function wrapped with @custom_gradient must " "be `ResourceVariable`s. Ensure that no `variable_scope` is created " "with `use_resource=False`.") # It is possible for the caller to pass in an input that is from a different # graph. Even though this is not valid we filter these out if they are not # from the output graph to make it easier for some code to migrate to custom # gradients. inputs = nest.flatten(args) outputs = nest.flatten(result) graphs = {getattr(o, "graph", None) for o in outputs} # Not all results may be tensors. However, we want to ensure that all outputs # are from the same graph and use that to filter the inputs. graphs.discard(None) # Discard non-graph outputs if graphs: if len(graphs) > 1: raise ValueError("All graph outputs should be from the same graph") output_graph = graphs.pop() filtered_inputs = [] for i in inputs: if i.graph != output_graph: logging.warn("%s does not belong to output graph %s", i, output_graph) else: filtered_inputs.append(i) inputs = filtered_inputs # The variables that grad_fn needs to return gradients for are the set of # variables used that are *not* part of the inputs. variables_in_tape = frozenset([ v.ref() for v in variable_watcher.watched_variables() ]) - frozenset(v.ref() for v in inputs) variables_in_subgraph = frozenset([ v.ref() for v in get_dependent_variables(input_ops=inputs, output_ops=outputs) ]) variables = list( [v.deref() for v in variables_in_subgraph.union(variables_in_tape)]) grad_argspec = tf_inspect.getfullargspec(grad_fn) variables_in_signature = ("variables" in grad_argspec.args or grad_argspec.varkw) if variables and not variables_in_signature: raise TypeError("If using @custom_gradient with a function that " "uses variables, then grad_fn must accept a keyword " "argument 'variables'.") if variables_in_signature and not variables: # User seems to intend to use variables but none were captured. logging.warn("@custom_gradient grad_fn has 'variables' in signature, but " "no ResourceVariables were used on the forward pass.") flat_result = nest.flatten(result) flat_result_len = len(flat_result) all_tensors = flat_result + inputs + variables def tape_grad_fn(*result_grads): """Custom grad fn wrapper.""" result_grads = result_grads[:flat_result_len] if variables: input_grads, variable_grads = grad_fn(*result_grads, variables=variables) if len(variable_grads) != len(variables): raise ValueError("Must return gradient for each variable from " "@custom_gradient grad_fn.") else: input_grads = grad_fn(*result_grads) variable_grads = [] # Need to return one value per input to the IdentityN, so pad the # gradients of the inputs of the custom_gradient function with the # gradients of the outputs as well. input_grads = nest.flatten(input_grads) return ([None] * flat_result_len) + input_grads + variable_grads @ops.RegisterGradient(name) def internal_grad_fn(unused_op, *result_grads): # pylint: disable=unused-variable """Custom grad fn wrapper.""" return tape_grad_fn(*result_grads) original_tensors = all_tensors with ops.get_default_graph().gradient_override_map({"IdentityN": name}): all_tensors = array_ops.identity_n(all_tensors) original_tensors = [ops.convert_to_tensor(x) for x in original_tensors] # Propagate handle data for happier shape inference for resource variables. for i, t in enumerate(original_tensors): if t.dtype == dtypes.resource and hasattr(t, "_handle_data"): all_tensors[i]._handle_data = t._handle_data # pylint: disable=protected-access tape_lib.record_operation( f.__name__, all_tensors, original_tensors, tape_grad_fn) for ot, t in zip(original_tensors, all_tensors): copy_handle_data(ot, t) return nest.pack_sequence_as( structure=result, flat_sequence=all_tensors[:flat_result_len])
def prune(self, feeds, fetches, name=None): name = name or "pruned" flat_feeds, flat_fetches = nest.flatten(feeds), nest.flatten(fetches) for f in flat_feeds: if not isinstance(f, ops.Tensor): raise ValueError("Feeds must be tensors.") # Ignoring all feeds that are captures allows prune to be called # using wrapped_func.inputs even when it uses variables internal_captures = self.graph.internal_captures flat_feeds = [f for f in flat_feeds if f not in internal_captures] tensor_fetches = [] operation_fetches = [] for f in flat_fetches: if isinstance(f, ops.Tensor): tensor_fetches.append(f) elif isinstance(f, ops.Operation): operation_fetches.append(f) else: raise ValueError("Fetches must be tensors or operations.") for f in flat_feeds + flat_fetches: if f.graph is not self._func_graph: raise ValueError( "Can only prune function whose feeds and fetches " "are from this graph (%s). Tensor %s from graph %s" % (self._func_graph, f, f.graph)) with self._func_graph.as_default(): pruned_graph = func_graph.FuncGraph(name) with ops.control_dependencies(operation_fetches): if tensor_fetches: identity_fetches = array_ops.identity_n(tensor_fetches) sink_tensor = identity_fetches[0] else: identity_fetches = [] sink_tensor = array_ops.zeros([]) lift_map = lift_to_graph.lift_to_graph([sink_tensor], pruned_graph, sources=flat_feeds + internal_captures) for original_fetch, identity_fetch in zip(tensor_fetches, identity_fetches): lift_map[original_fetch] = lift_map[identity_fetch] pruned_graph.outputs.extend(lift_map[x] for x in flat_fetches if isinstance(x, ops.Tensor)) pruned_graph.control_outputs.extend( [lift_map[operation] for operation in operation_fetches]) if not tensor_fetches: pruned_graph.outputs.append(lift_map[sink_tensor]) for external_capture, internal_capture in self.graph.captures.items(): pruned_graph.captures[external_capture] = lift_map[ internal_capture] pruned_graph.inputs.extend(lift_map[x] for x in flat_feeds) pruned_graph.inputs.extend(pruned_graph.captures.values()) pruned_graph.variables = self.graph.variables def _structured_output_mapping(fetched): lifted = lift_map[fetched] if isinstance(lifted, ops.Operation): return None return lifted pruned_graph.structured_outputs = nest.map_structure( _structured_output_mapping, fetches) pruned_fn = WrappedFunction(pruned_graph, variable_holder=self._variable_holder) pruned_fn._num_positional_args = len(flat_feeds) # pylint: disable=protected-access pruned_fn._arg_keywords = [] # pylint: disable=protected-access return pruned_fn
def prune(self, feeds, fetches): flat_feeds, flat_fetches = nest.flatten(feeds), nest.flatten(fetches) for f in flat_feeds: if not isinstance(f, ops.Tensor): raise ValueError("Feeds must be tensors.") # Ignoring all feeds that are captures allows prune to be called # using wrapped_func.inputs even when it uses variables internal_captures = self.graph.internal_captures flat_feeds = [f for f in flat_feeds if f not in internal_captures] tensor_fetches = [] operation_fetches = [] for f in flat_fetches: if isinstance(f, ops.Tensor): tensor_fetches.append(f) elif isinstance(f, ops.Operation): operation_fetches.append(f) else: raise ValueError("Fetches must be tensors or operations.") for f in flat_feeds + flat_fetches: if f.graph is not self._func_graph: raise ValueError( "Can only prune function whose feeds and fetches " "are from this graph (%s). Tensor %s from graph %s" % ( self._func_graph, f, f.graph)) with self._func_graph.as_default(): pruned_graph = func_graph.FuncGraph("pruned") with ops.control_dependencies(operation_fetches): if tensor_fetches: identity_fetches = array_ops.identity_n(tensor_fetches) sink_tensor = identity_fetches[0] else: identity_fetches = [] sink_tensor = array_ops.zeros([]) lift_map = lift_to_graph.lift_to_graph( [sink_tensor], pruned_graph, sources=flat_feeds + internal_captures) for original_fetch, identity_fetch in zip( tensor_fetches, identity_fetches): lift_map[original_fetch] = lift_map[identity_fetch] pruned_graph.outputs.extend( lift_map[x] for x in flat_fetches if isinstance(x, ops.Tensor)) if not tensor_fetches: pruned_graph.outputs.append(lift_map[sink_tensor]) for external_capture, internal_capture in self.graph.captures.items(): pruned_graph.captures[external_capture] = lift_map[internal_capture] pruned_graph.inputs.extend(lift_map[x] for x in flat_feeds) pruned_graph.inputs.extend(pruned_graph.captures.values()) pruned_graph.variables = self.graph.variables def _structured_output_mapping(fetched): lifted = lift_map[fetched] if isinstance(lifted, ops.Operation): return None return lifted pruned_graph.structured_outputs = nest.map_structure( _structured_output_mapping, fetches) pruned_fn = WrappedFunction( pruned_graph, variable_holder=self._variable_holder) pruned_fn._num_positional_args = len(flat_feeds) # pylint: disable=protected-access pruned_fn._arg_keywords = [] # pylint: disable=protected-access return pruned_fn
def testString(self): source = [b"A", b"b", b"C", b"d", b"E", b"f"] [value] = self.evaluate(array_ops.identity_n([source])) self.assertAllEqual(source, value)
def testString(self): source = [b"A", b"b", b"C", b"d", b"E", b"f"] with self.test_session() as sess: [value] = sess.run(array_ops.identity_n([source])) self.assertAllEqual(source, value)
def testMixedTypeListInputEagerFallbackDifferentArity(self): array_ops.identity_n([1, 1]) array_ops.identity_n([1, 1, 1])
def testMixedTypeListInputFastPathDifferentArity(self): # This tests that the FunctionDef cache key contains the number of args. array_ops.identity_n([self._m_2_by_2, self._m_2_by_2]) array_ops.identity_n([self._m_2_by_2, self._m_2_by_2, self._m_2_by_2])
def _graph_mode_decorator(f, args, kwargs): """Implement custom gradient decorator for graph mode.""" # TODO(rsepassi): Add support for kwargs if kwargs: raise ValueError( "The custom_gradient decorator currently supports keywords " "arguments only when eager execution is enabled.") name = "CustomGradient-%s" % ops.uid() default_graph = ops.get_default_graph() def convert_arg(x): x = ops.convert_to_tensor(x) # If graph building, be sure to capture all inputs if default_graph.building_function and x.graph != default_graph: x = default_graph.capture(x) return x args = nest.map_structure(convert_arg, args) # Checking global and local variables attempts to ensure that no non-resource # Variables are added to the graph. current_var_scope = variable_scope.get_variable_scope() before_vars = set([ v.ref() for v in current_var_scope.global_variables() + current_var_scope.local_variables() ]) with tape_lib.VariableWatcher() as variable_watcher: result, grad_fn = f(*args) args = nest.flatten(args) flat_result = nest.flatten(result) flat_result_len = len(flat_result) after_vars = set([ v.ref() for v in current_var_scope.global_variables() + current_var_scope.local_variables() ]) new_vars = after_vars - before_vars new_vars_list = [v.deref() for v in new_vars] for v in new_vars_list: if not resource_variable_ops.is_resource_variable(v): raise TypeError( "All variables used by a function wrapped with @custom_gradient must " "be `ResourceVariable`s. Ensure that no `variable_scope` is created " "with `use_resource=False`.") # The variables that grad_fn needs to return gradients for are the set of # variables used that are *not* part of the inputs. variables_in_tape = frozenset( [v.ref() for v in variable_watcher.watched_variables()]) variables_in_subgraph = frozenset([ v.ref() for v in _get_dependent_variables(input_ops=args, output_ops=flat_result) ]) variables = list( [v.deref() for v in variables_in_subgraph.union(variables_in_tape)]) grad_argspec = tf_inspect.getfullargspec(grad_fn) variables_in_signature = ("variables" in grad_argspec.args or "variables" in grad_argspec.kwonlyargs or grad_argspec.varkw) if variables and not variables_in_signature: raise TypeError( "@tf.custom_gradient grad_fn must accept keyword argument 'variables', " "since function uses variables: {}".format(variables)) if variables_in_signature and not variables: # User seems to intend to use variables but none were captured. logging.warn( "@custom_gradient grad_fn has 'variables' in signature, but " "no ResourceVariables were used on the forward pass.") all_tensors = flat_result + args + variables def tape_grad_fn(*result_grads): """Custom grad fn wrapper.""" result_grads = result_grads[:flat_result_len] if variables: input_grads, variable_grads = grad_fn(*result_grads, variables=variables) if len(variable_grads) != len(variables): raise ValueError("Must return gradient for each variable from " "@custom_gradient grad_fn.") else: input_grads = grad_fn(*result_grads) variable_grads = [] # Need to return one value per input to the IdentityN, so pad the # gradients of the inputs of the custom_gradient function with the # gradients of the outputs as well. input_grads = nest.flatten(input_grads) return ([None] * flat_result_len) + input_grads + variable_grads @ops.RegisterGradient(name) def internal_grad_fn(unused_op, *result_grads): # pylint: disable=unused-variable """Custom grad fn wrapper.""" return tape_grad_fn(*result_grads) original_tensors = all_tensors with ops.get_default_graph().gradient_override_map({"IdentityN": name}): all_tensors = array_ops.identity_n(all_tensors) original_tensors = [ops.convert_to_tensor(x) for x in original_tensors] # Propagate handle data for happier shape inference for resource variables. for i, t in enumerate(original_tensors): if t.dtype == dtypes.resource and hasattr(t, "_handle_data"): all_tensors[i]._handle_data = t._handle_data # pylint: disable=protected-access tape_lib.record_operation(f.__name__, all_tensors, original_tensors, tape_grad_fn) for ot, t in zip(original_tensors, all_tensors): copy_handle_data(ot, t) return nest.pack_sequence_as(structure=result, flat_sequence=all_tensors[:flat_result_len])
def loop_fn(i): return array_ops.identity_n([x, array_ops.gather(x, i)])
def testString(self): source = [b"A", b"b", b"C", b"d", b"E", b"f"] with self.cached_session() as sess: [value] = sess.run(array_ops.identity_n([source])) self.assertAllEqual(source, value)
def testMixedTypeListInputFastPath(self): array_ops.identity_n([self._m_2_by_2, self._m_2_by_2])
def _graph_mode_decorator(f, *args, **kwargs): """Implement custom gradient decorator for graph mode.""" # TODO(rsepassi): Add support for kwargs if kwargs: raise ValueError( "The custom_gradient decorator currently supports keywords " "arguments only when eager execution is enabled.") name = "CustomGradient-%s" % ops.uid() args = [ops.convert_to_tensor(x) for x in args] # Checking global and local variables attempts to ensure that no non-resource # Variables are added to the graph. current_var_scope = variable_scope.get_variable_scope() before_vars = set(current_var_scope.global_variables() + current_var_scope.local_variables()) with backprop.GradientTape() as tape: result, grad_fn = f(*args) after_vars = set(current_var_scope.global_variables() + current_var_scope.local_variables()) new_vars = after_vars - before_vars for v in new_vars: if not isinstance(v, resource_variable_ops.ResourceVariable): raise TypeError( "All variables used by a function wrapped with @custom_gradient must " "be `ResourceVariable`s. Ensure that no `variable_scope` is created " "with `use_resource=False`.") # The variables that grad_fn needs to return gradients for are the set of # variables used that are *not* part of the inputs. variables = list(set(tape.watched_variables()) - set(args)) grad_argspec = tf_inspect.getfullargspec(grad_fn) variables_in_signature = ("variables" in grad_argspec.args or grad_argspec.varkw) if variables and not variables_in_signature: raise TypeError("If using @custom_gradient with a function that " "uses variables, then grad_fn must accept a keyword " "argument 'variables'.") if variables_in_signature and not variables: # User seems to intend to use variables but none were captured. if not variable_scope.get_variable_scope().use_resource: raise TypeError("If using @custom_gradient with a function that " "uses variables, the enclosing variable scope must " "have use_resource=True.") else: logging.warn("@custom_gradient grad_fn has 'variables' in signature, but " "no ResourceVariables were used on the forward pass.") flat_result = nest.flatten(result) all_tensors = flat_result + args + variables def tape_grad_fn(*result_grads): """Custom grad fn wrapper.""" result_grads = result_grads[:len(flat_result)] if variables: input_grads, variable_grads = grad_fn(*result_grads, variables=variables) if len(variable_grads) != len(variables): raise ValueError("Must return gradient for each variable from " "@custom_gradient grad_fn.") else: input_grads = grad_fn(*result_grads) variable_grads = [] # Need to return one value per input to the IdentityN, so pad the # gradients of the inputs of the custom_gradient function with the # gradients of the outputs as well. input_grads = nest.flatten(input_grads) return ([None] * len(flat_result)) + input_grads + variable_grads @ops.RegisterGradient(name) def internal_grad_fn(unused_op, *result_grads): # pylint: disable=unused-variable """Custom grad fn wrapper.""" return tape_grad_fn(*result_grads) original_tensors = all_tensors with ops.get_default_graph().gradient_override_map({"IdentityN": name}): all_tensors = array_ops.identity_n(all_tensors) # Propagate handle data for happier shape inference for resource variables. for i, t in enumerate(original_tensors): if t.dtype == dtypes.resource and hasattr(t, "_handle_data"): all_tensors[i]._handle_data = t._handle_data # pylint: disable=protected-access tape_lib.record_operation( f.__name__, all_tensors, original_tensors, tape_grad_fn) for ot, t in zip(original_tensors, all_tensors): copy_handle_data(ot, t) return nest.pack_sequence_as( structure=result, flat_sequence=all_tensors[:len(flat_result)])
def testMixedTypeListInputEagerFallback(self): array_ops.identity_n([1, 1])
def loop_fn(i): return array_ops.identity_n([x, array_ops.gather(x, i)])
def testListInputOutput(self): self.skipTest("b/185403393") array_ops.identity_n([self._m_2_by_2, self._m_2_by_2])