def test_works_with_registered(self): class CustomClass(object): def value(self): return ops.convert_to_tensor(42.) ops.register_tensor_conversion_function( CustomClass, lambda value, **_: value.value()) tf_utils.register_symbolic_tensor_type(CustomClass) if context.executing_eagerly(): self.assertFalse(tf_utils.is_symbolic_tensor( variables.Variable(name='blah', initial_value=0.))) self.assertFalse(tf_utils.is_symbolic_tensor( ops.convert_to_tensor(0.))) self.assertFalse(tf_utils.is_symbolic_tensor( sparse_tensor.SparseTensor( indices=[[0, 0], [1, 2]], values=[1, 2], dense_shape=[3, 4]))) self.assertFalse(tf_utils.is_symbolic_tensor(CustomClass())) else: self.assertTrue(tf_utils.is_symbolic_tensor( variables.Variable(name='blah', initial_value=0.))) self.assertTrue(tf_utils.is_symbolic_tensor( ops.convert_to_tensor(0.))) self.assertTrue(tf_utils.is_symbolic_tensor( sparse_tensor.SparseTensor( indices=[[0, 0], [1, 2]], values=[1, 2], dense_shape=[3, 4]))) self.assertTrue(tf_utils.is_symbolic_tensor(CustomClass()))
def test_enables_nontensor_plumbing(self): # Setup. class Foo(object): def __init__(self, input_): self._input = input_ self.value = ops.convert_to_tensor(42.) ops.register_tensor_conversion_function( Foo, lambda x, *args, **kwargs: x.value) tf_utils.register_symbolic_tensor_type(Foo) class PlumbingLayer(keras.layers.Lambda): def __init__(self, fn, **kwargs): def _fn(*fargs, **fkwargs): d = fn(*fargs, **fkwargs) x = ops.convert_to_tensor(d) d.shape = x.shape d.get_shape = x.get_shape return d, x super(PlumbingLayer, self).__init__(_fn, **kwargs) self._enter_dunder_call = False def __call__(self, inputs, *args, **kwargs): self._enter_dunder_call = True d, _ = super(PlumbingLayer, self).__call__(inputs, *args, **kwargs) self._enter_dunder_call = False return d def call(self, inputs, *args, **kwargs): d, v = super(PlumbingLayer, self).call(inputs, *args, **kwargs) if self._enter_dunder_call: return d, v return d # User-land. model = keras.Sequential([ keras.layers.InputLayer([]), PlumbingLayer(Foo), # Makes a `Foo` object. ]) # Let's ensure Keras graph history is preserved by composing the models. model = keras.Model(model.inputs, model(model.outputs)) # Now we instantiate the model and verify we have a `Foo` object, not a # `Tensor`. y = model(ops.convert_to_tensor(7.)) self.assertIsInstance(y, Foo)
def testFullDelegationControlUsingRegistry(self): class NumpyArraySubclass(np.ndarray): def __radd__(self, lhs): return "Works!" def raise_to_delegate(value, dtype=None, name=None, as_ref=False): del value, dtype, name, as_ref # Unused. raise TypeError ops.register_tensor_conversion_function(NumpyArraySubclass, raise_to_delegate, priority=0) tensor = ops.convert_to_tensor([[10.0, 20.0]]) rhs = NumpyArraySubclass(shape=(1, 2), buffer=np.array([1.0, 2.0])) res = tensor + rhs self.assertEqual(res, "Works!")
def register_tensor_conversion(convertable, name=None, overload_operators=True, priority=10): # higher than any tf conversion def _dense_var_to_tensor(var, dtype=None, name=None, as_ref=False): return var._dense_var_to_tensor(dtype=dtype, name=name, as_ref=as_ref) ops.register_tensor_conversion_function(convertable, _dense_var_to_tensor, priority=priority) if name: pass # _pywrap_utils.RegisterType(name, convertable) if overload_operators: convertable._OverloadAllOperators()
def testFullDelegationControlUsingRegistry(self): class NumpyArraySubclass(np.ndarray): def __radd__(self, lhs): return "Works!" def raise_to_delegate(value, dtype=None, name=None, as_ref=False): del value, dtype, name, as_ref # Unused. raise TypeError ops.register_tensor_conversion_function( NumpyArraySubclass, raise_to_delegate, priority=0) tensor = ops.convert_to_tensor([[10.0, 20.0]]) rhs = NumpyArraySubclass(shape=(1, 2), buffer=np.array([1.0, 2.0])) res = tensor + rhs self.assertEqual(res, "Works!")
def test_works_with_registered(self): class CustomClass(object): def value(self): return ops.convert_to_tensor_v2_with_dispatch(42.) ops.register_tensor_conversion_function( CustomClass, lambda value, **_: value.value()) tf_utils.register_symbolic_tensor_type(CustomClass) if context.executing_eagerly(): self.assertFalse( tf_utils.is_symbolic_tensor( variables.Variable(name='blah', initial_value=0.))) self.assertFalse( tf_utils.is_symbolic_tensor( ops.convert_to_tensor_v2_with_dispatch(0.))) self.assertFalse( tf_utils.is_symbolic_tensor( sparse_tensor.SparseTensor(indices=[[0, 0], [1, 2]], values=[1, 2], dense_shape=[3, 4]))) self.assertFalse(tf_utils.is_symbolic_tensor(CustomClass())) else: self.assertTrue( tf_utils.is_symbolic_tensor( variables.Variable(name='blah', initial_value=0.))) self.assertTrue( tf_utils.is_symbolic_tensor( ops.convert_to_tensor_v2_with_dispatch(0.))) self.assertTrue( tf_utils.is_symbolic_tensor( sparse_tensor.SparseTensor(indices=[[0, 0], [1, 2]], values=[1, 2], dense_shape=[3, 4]))) self.assertTrue(tf_utils.is_symbolic_tensor(CustomClass()))
return self._read_variable_op() def _read_variable_op(self): with ops.control_dependencies([self._parent_op]): return gen_resource_variable_ops.read_variable_op(self._handle, self._dtype) def set_shape(self, shape): self._shape = shape @property def op(self): """The op for this variable.""" return self._parent_op ops.register_tensor_conversion_function(_UnreadVariable, _dense_var_to_tensor) ops.register_dense_tensor_like_type(_UnreadVariable) # Register a conversion function which reads the value of the variable, # allowing instances of the class to be used as tensors. # Note: registering for Variable after ResourceVariable because inheritance will # otherwise lead to the wrong behavior. ops.register_tensor_conversion_function(ResourceVariable, _dense_var_to_tensor) ops.register_tensor_conversion_function( variables.Variable, variables.Variable._TensorConversionFunction) # pylint: disable=protected-access # pylint: disable=protected-access ResourceVariable._OverloadAllOperators() ops.register_dense_tensor_like_type(ResourceVariable)
collections: any collections in which this operation should be inserted. trainable: whether this read is to be used for training. Returns: the read operation. """ with ops.name_scope("Read"): value = gen_resource_variable_ops.read_variable_op( self._handle, dtype=self._dtype) _register_variable_read(value, collections=collections, trainable=trainable) return value def sparse_read(self, indices, collections=None, trainable=True, name=None): with ops.name_scope("Gather" if name is None else name): value = gen_resource_variable_ops.resource_gather( self._handle, indices, dtype=self._dtype) _register_variable_read(value, collections=collections, trainable=trainable) return value # pylint: disable=unused-argument def _dense_var_to_tensor(var, dtype=None, name=None, as_ref=False): if dtype is not None and dtype != var.value.dtype: print("trying to switch the dtype to ", dtype, " from ", var.value.dtype) return NotImplemented return var.value # Register a conversion function which reads the value of the variable, # allowing instances of the class to be used as tensors. ops.register_tensor_conversion_function(ResourceVariable, _dense_var_to_tensor)
@property def shape(self): return self.data.shape @property def dtype(self): return self.data.dtype def fail_on_convert(x, **kwargs): _ = x _ = kwargs raise TypeError('Cannot convert DummyArrayLike to a tensor') ops.register_tensor_conversion_function(DummyArrayLike, fail_on_convert) class DataAdapterTestBase(keras_parameterized.TestCase): def setUp(self): super(DataAdapterTestBase, self).setUp() self.batch_size = 5 self.numpy_input = np.zeros((50, 10)) self.numpy_target = np.ones(50) self.tensor_input = constant_op.constant(2.0, shape=(50, 10)) self.tensor_target = array_ops.ones((50, )) self.arraylike_input = DummyArrayLike(self.numpy_input) self.arraylike_target = DummyArrayLike(self.numpy_target) self.dataset_input = dataset_ops.DatasetV2.from_tensor_slices( (self.numpy_input, self.numpy_target)).shuffle(50).batch(self.batch_size)
if not isinstance(self.data, ops.EagerTensor): raise TypeError('Indexing using symbolic tensor is not allowed') return self.data.numpy().item() def tolist(self): return self.data.numpy().tolist() def __str__(self): return 'ndarray<{}>'.format(self.data.__str__()) def __repr__(self): return 'ndarray<{}>'.format(self.data.__repr__()) def tensor_to_ndarray(tensor): return ndarray.from_tensor(tensor) def ndarray_to_tensor(arr, dtype=None, name=None, as_ref=False): if as_ref: raise ValueError('as_ref is not supported.') if dtype and dtypes.as_dtype(arr.dtype) != dtype: return math_ops.cast(arr.data, dtype) result_t = arr.data if name: result_t = array_ops.identity(result_t, name=name) return result_t ops.register_tensor_conversion_function(ndarray, ndarray_to_tensor)
def _map_resources(self, save_options): """For implementing `Trackable`.""" # By delegating this method to the wrapped variable, SavedModel with # AggregatingVariable are identical to SavedModel with normal variables. obj_map, resource_map = self._v._map_resources(save_options) # pylint:disable=protected-access obj_map[self] = obj_map[self._v] return obj_map, resource_map # Register a conversion function which reads the value of the variable, # allowing instances of the class to be used as tensors. def _tensor_conversion_aggregate(var, dtype=None, name=None, as_ref=False): return var._dense_var_to_tensor(dtype, name, as_ref) # pylint: disable=protected-access ops.register_tensor_conversion_function(AggregatingVariable, _tensor_conversion_aggregate) # Register a conversion function which reads the value of the variable, # allowing instances of the class to be used as tensors. def _tensor_conversion_caching(var, dtype=None, name=None, as_ref=False): return var._dense_var_to_tensor(dtype, name, as_ref) # pylint: disable=protected-access ops.register_tensor_conversion_function(CachingVariable, _tensor_conversion_caching) CachingVariable._overload_overloadable_operators() # pylint: disable=protected-access class DistributedTable(lookup_ops.StaticHashTable):
dense_shape_value = tensor_util.constant_value(value.dense_shape) if dense_shape_value is not None: num_elements = np.prod(dense_shape_value) if num_elements >= _LARGE_SPARSE_NUM_ELEMENTS: warnings.warn( "Converting sparse IndexedSlices to a dense Tensor with %d elements. " "This may consume a large amount of memory." % num_elements) else: warnings.warn( "Converting sparse IndexedSlices to a dense Tensor of unknown shape. " "This may consume a large amount of memory.") return math_ops.unsorted_segment_sum( value.values, value.indices, value.dense_shape[0], name=name) ops.register_tensor_conversion_function(ops.IndexedSlices, _IndexedSlicesToTensor) def _MarkReachedOps(from_ops, reached_ops): """Mark all ops reached from "from_ops". Args: from_ops: list of Operations. reached_ops: list of booleans, indexed by operation id. """ queue = collections.deque() queue.extend(from_ops) while queue: op = queue.popleft() if not reached_ops[op._id]: reached_ops[op._id] = True
return self.read_value().__matmul__(o) except AttributeError: # See https://docs.python.org/3/library/constants.html#NotImplemented return NotImplemented def __rmatmul__(self, o): try: return self.read_value().__rmatmul__(o) except AttributeError: # See https://docs.python.org/3/library/constants.html#NotImplemented return NotImplemented # pylint: enable=multiple-statements ops.register_tensor_conversion_function(AutoCastVariable, AutoCastVariable._dense_var_to_tensor) # pylint:disable=protected-access def create_autocast_variable(variable): """Creates an AutoCastVariable that wraps another variable. This typically just returns `AutoCastVariable(variable)`. But, if the variable is a DistributedVariable or one of its subclasses, we instead dynamically create a class that subclasses from both AutoCastVariable and variable.__class__. This is so the returned variable will still pass `isinstance(variable, variable.__class__)`, which is required for DistributedVariables and its subclasses to work properly. Args: variable: A floating-point resource variable to wrap.
def test_enables_nontensor_plumbing(self): # Setup. class Foo(object): def __init__(self, input_): self._input = input_ self.value = ops.convert_to_tensor([[42.]]) @property def dtype(self): return self.value.dtype ops.register_tensor_conversion_function( Foo, lambda x, *args, **kwargs: x.value) tf_utils.register_symbolic_tensor_type(Foo) class PlumbingLayer(keras.layers.Lambda): def __init__(self, fn, **kwargs): def _fn(*fargs, **fkwargs): d = fn(*fargs, **fkwargs) x = ops.convert_to_tensor(d) d.shape = x.shape d.get_shape = x.get_shape return d, x super(PlumbingLayer, self).__init__(_fn, **kwargs) self._enter_dunder_call = False def __call__(self, inputs, *args, **kwargs): self._enter_dunder_call = True d, _ = super(PlumbingLayer, self).__call__(inputs, *args, **kwargs) self._enter_dunder_call = False return d def call(self, inputs, *args, **kwargs): d, v = super(PlumbingLayer, self).call(inputs, *args, **kwargs) if self._enter_dunder_call: return d, v return d # User-land. model = keras.Sequential([ keras.layers.InputLayer((1, )), PlumbingLayer(Foo), # Makes a `Foo` object. ]) # Let's ensure Keras graph history is preserved by composing the models. model = keras.Model(model.inputs, model(model.outputs)) # Now we instantiate the model and verify we have a `Foo` object, not a # `Tensor`. y = model(ops.convert_to_tensor([[7.]])) self.assertIsInstance(y, Foo) # Confirm that (custom) loss sees `Foo` instance, not Tensor. obtained_prediction_box = [None] def custom_loss(y_obs, y_pred): del y_obs obtained_prediction_box[0] = y_pred return y_pred # Apparently `compile` calls the loss function enough to trigger the # side-effect. model.compile('SGD', loss=custom_loss) self.assertIsInstance(obtained_prediction_box[0], Foo)
self.sess.run(self.n_samples.initializer) if n: if not isinstance(self.n_samples, tf.Variable): raise RuntimeError( "Cannot set a new `n` if not a Tensor-like object was given" ) self.n_samples.load(value=n, session=self.sess) self.sess.run(self.sample_holder.initializer) self._initial_resampled = True def _dense_var_to_tensor(var, dtype=None, name=None, as_ref=False): return var._dense_var_to_tensor(dtype=dtype, name=name, as_ref=as_ref) ops.register_tensor_conversion_function(Data, _dense_var_to_tensor) fetch_function = lambda data: ([data.value()], lambda val: val[0]) feed_function = lambda data, feed_val: [(data.value(), feed_val)] feed_function_for_partial_run = lambda data: [data.value()] from tensorflow.python.client.session import register_session_run_conversion_functions # ops.register_dense_tensor_like_type() register_session_run_conversion_functions( tensor_type=Data, fetch_function=fetch_function, feed_function=feed_function, feed_function_for_partial_run=feed_function_for_partial_run) Data._OverloadAllOperators()
return self.read_value() # Register a conversion function which reads the value of the variable, # allowing instances of the class to be used as tensors. def _tensor_conversion(var, dtype=None, name=None, as_ref=False): return var._dense_var_to_tensor(dtype=dtype, name=name, as_ref=as_ref) # pylint: disable=protected-access def replicated_fetch_function(var): # pylint: disable=protected-access return ([var._dense_var_to_tensor()], lambda v: v[0]) # pylint: enable=protected-access ops.register_tensor_conversion_function(ReplicatedVariable, _tensor_conversion) ops.register_dense_tensor_like_type(ReplicatedVariable) session_lib.register_session_run_conversion_functions( ReplicatedVariable, replicated_fetch_function) def replicated_scope(num_replicas): """Variable scope for constructing replicated variables.""" def _replicated_variable_getter(getter, name, *args, **kwargs): """Getter that constructs replicated variables.""" collections = kwargs.pop("collections", None) if collections is None: collections = [ops.GraphKeys.GLOBAL_VARIABLES] kwargs["collections"] = []
def _saveable_factory(name=self._common_name): return _MirroredSaveable(self, self._primary_var, name) return {checkpointable.VARIABLE_VALUE_KEY: _saveable_factory} # Register a conversion function which reads the value of the variable, # allowing instances of the class to be used as tensors. def _tensor_conversion_mirrored(var, dtype=None, name=None, as_ref=False): # Try to avoid assignments to and other mutations of MirroredVariable # state except through a DistributionStrategy.update() call. assert not as_ref return ops.internal_convert_to_tensor( var.get(), dtype=dtype, name=name, as_ref=as_ref) ops.register_tensor_conversion_function(MirroredVariable, _tensor_conversion_mirrored) class _TowerLocalSaveable(saver.BaseSaverBuilder.SaveableObject): """Class for defining how to restore a TowerLocalVariable.""" def __init__(self, tower_local_variable, name): self._tower_local_variable = tower_local_variable # We use a callable so that we don't have to evaluate this expression # in the case where we are trying to restore instead of save. def tensor(): return distribute_lib.get_distribution_strategy().read_var( tower_local_variable) spec = saver.BaseSaverBuilder.SaveSpec( tensor=tensor, slice_spec="",
captured_value = tensor_map.get(ops.tensor_id(value), None) if captured_value is None: captured_value = graph_placeholder( dtype=dtype or value.dtype, shape=value.shape, name=name) if captured_value.dtype == dtypes.resource: captured_value._handle_data = value._handle_data # pylint: disable=protected-access tensor_map[ops.tensor_id(value)] = (value, captured_value) else: captured_value = captured_value[1] return captured_value # TODO(apassos): it'd be really nice if we could scope this registration. # Note that we register this at a higher priority than ops.Tensor since we want # to handle subclass specific conversion before a superclass conversion. ops.register_tensor_conversion_function( tensor.Tensor, _convert_to_graph_constant, priority=-1) class _CapturingContext(object): """Tracks references to Tensors outside this context while it is active.""" def __init__(self): # known_ops are ops which are created while this context is active self.known_ops = set() # captured_tensors are all tensors referenced to by ops in this context but # not produced in it self.captured_tensors = set() def AddOp(self, op): # pylint: disable=invalid-name if op.type in ["Variable", "VariableV2", "VarHandleOp"]:
return property(Cache().__getitem__) ``` However, that implementation holds class instances as keys, and as a result blocks garbage collection. (And modifying it to use weakref's as keys raises the lookup overhead to ~0.4 us) As a result, the WeakKeyDictionary implementation below turns out to be more prudent. Args: f: The function to cache. Returns: f decorated with simple caching behavior. """ cache = weakref.WeakKeyDictionary() @functools.wraps(f) def wrapped(item): output = cache.get(item) if output is None: cache[item] = output = f(item) return output return wrapped ops.register_tensor_conversion_function( TrackableAsset, lambda asset, **kw: ops.internal_convert_to_tensor(asset.asset_path, **kw))
container = "" return gen_ev_ops.ev_handle_op(shape=shape, shared_name=shared_name, name=name, Tkey=self._ktype, Tvalue=dtype, container=container) # Register a conversion function which reads the value of the variable, # allowing instances of the class to be used as tensors. # Note: registering for Variable after EmbeddingVariable because inheritance will # otherwise lead to the wrong behavior. #ops.register_tensor_conversion_function(EmbeddingVariable, _dense_var_to_tensor) ops.register_tensor_conversion_function( variables.Variable, variables.Variable._TensorConversionFunction) # pylint: disable=protected-access @ops.RegisterGradient("EVGather") def _GatherGrad(op, grad): """Gradient for gather op.""" handle = op.inputs[0] indices = op.inputs[1] params_shape = gen_ev_ops.ev_shape(handle, Tkey=indices.dtype, Tvalue=grad.dtype) size = array_ops.expand_dims(array_ops.size(indices), 0) values_shape = array_ops.concat([size, params_shape[1:]], 0) values = array_ops.reshape(grad, values_shape) indices = array_ops.reshape(indices, size) return [ops.IndexedSlices(values, indices, params_shape), None, None]
var_list.append(op.outputs[0]) if not var_list: return None else: ranks = [] for var in var_list: with ops.device(var.device): ranks.append(array_ops.rank(var)) if len(ranks) == 1: return ranks[0] else: return array_ops.pack(ranks) # pylint: disable=protected-access ops.register_tensor_conversion_function(Variable, Variable._TensorConversionFunction) Variable._OverloadAllOperators() # pylint: enable=protected-access ops.register_proto_function(ops.GraphKeys.VARIABLES, proto_type=variable_pb2.VariableDef, to_proto=Variable.to_proto, from_proto=Variable.from_proto) ops.register_proto_function(ops.GraphKeys.TRAINABLE_VARIABLES, proto_type=variable_pb2.VariableDef, to_proto=Variable.to_proto, from_proto=Variable.from_proto) ops.register_proto_function(ops.GraphKeys.MOVING_AVERAGE_VARIABLES, proto_type=variable_pb2.VariableDef, to_proto=Variable.to_proto, from_proto=Variable.from_proto)
delta_get_op = delta_staging_area.get()[0] # Return the actual updates. The colocation constraint will be reapplied. return self.real_var.assign_sub(delta_get_op, read_value=read_value) @staticmethod # pylint: disable=bad-staticmethod-argument,invalid-name def _TensorConversionFunction(self, dtype=None, name=None, as_ref=False): """Utility function for converting a StagedModelVariable to a Tensor.""" del dtype, name # unused: this function returns the cached ref or value. if as_ref: return self._ref() else: return self._value() ops.register_tensor_conversion_function( StagedModelVariable, StagedModelVariable._TensorConversionFunction) # pylint: disable=protected-access class StagedVariableGetter(object): """A variable getter through staging buffers on devices. Instead of a caching device, this getter tracks where the variable is used. And on each device, it goes through a staging buffer. """ def __init__(self, device_num, devices, cpu_device, variable_mgr): """Initializer for StagedVariableGetter. Args: device_num: the current device index. devices: a list of all the devices to build towers. cpu_device: a cpu_device for this replica. If None, no cpu-caching is
# to fail. for i, inp in enumerate(inputs): if inp.graph is not self: inputs[i] = capture_value(self.captures, inp, inp.dtype, inp.op.name) return super(CapturingGraph, self).create_op(op_type, inputs, dtypes, input_types, name, attrs, op_def, compute_shapes, compute_device) # TODO(apassos): it'd be really nice if we could scope this registration. # Note that we register this at a higher priority than ops.Tensor since we want # to handle subclass specific conversion before a superclass conversion. ops.register_tensor_conversion_function(ops.EagerTensor, _convert_to_graph_tensor, priority=-1) # pylint: disable=invalid-name class HelperContext(object): """ControlFlowContext with a customizable AddOp method.""" def __init__(self, add_op_internal): self._add_op_internal = add_op_internal self._values = set() # control flow code sometimes updates this. def _AddOpInternal(self, op): self._add_op_internal(op) @property def outer_context(self):
compute_device=True): # TODO(apassos) this should do some form of alias analysis as ops which # forward the resources such as Identity and Switch can cause serialization # to fail. for i, inp in enumerate(inputs): if inp.graph is not self: inputs[i] = capture_value(self.captures, inp, inp.dtype, inp.op.name) return super(CapturingGraph, self).create_op( op_type, inputs, dtypes, input_types, name, attrs, op_def, compute_shapes, compute_device) # TODO(apassos): it'd be really nice if we could scope this registration. # Note that we register this at a higher priority than ops.Tensor since we want # to handle subclass specific conversion before a superclass conversion. ops.register_tensor_conversion_function( ops.EagerTensor, _convert_to_graph_tensor, priority=-1) # pylint: disable=invalid-name class HelperContext(object): """ControlFlowContext with a customizable AddOp method.""" def __init__(self, add_op_internal): self._add_op_internal = add_op_internal self._values = set() # control flow code sometimes updates this. def _AddOpInternal(self, op): self._add_op_internal(op) @property def outer_context(self):
def test_enables_nontensor_plumbing(self): # Setup. class Foo(object): def __init__(self, input_): self._input = input_ self.value = ops.convert_to_tensor(42.) @property def dtype(self): return self.value.dtype ops.register_tensor_conversion_function( Foo, lambda x, *args, **kwargs: x.value) tf_utils.register_symbolic_tensor_type(Foo) class PlumbingLayer(keras.layers.Lambda): def __init__(self, fn, **kwargs): def _fn(*fargs, **fkwargs): d = fn(*fargs, **fkwargs) x = ops.convert_to_tensor(d) d.shape = x.shape d.get_shape = x.get_shape return d, x super(PlumbingLayer, self).__init__(_fn, **kwargs) self._enter_dunder_call = False def __call__(self, inputs, *args, **kwargs): self._enter_dunder_call = True d, _ = super(PlumbingLayer, self).__call__(inputs, *args, **kwargs) self._enter_dunder_call = False return d def call(self, inputs, *args, **kwargs): d, v = super(PlumbingLayer, self).call(inputs, *args, **kwargs) if self._enter_dunder_call: return d, v return d # User-land. model = keras.Sequential([ keras.layers.InputLayer([]), PlumbingLayer(Foo), # Makes a `Foo` object. ]) # Let's ensure Keras graph history is preserved by composing the models. model = keras.Model(model.inputs, model(model.outputs)) # Now we instantiate the model and verify we have a `Foo` object, not a # `Tensor`. y = model(ops.convert_to_tensor(7.)) self.assertIsInstance(y, Foo) # Confirm that (custom) loss sees `Foo` instance, not Tensor. obtained_prediction_box = [None] def custom_loss(y_obs, y_pred): del y_obs obtained_prediction_box[0] = y_pred return y_pred # Apparently `compile` calls the loss function enough to trigger the # side-effect. model.compile('SGD', loss=custom_loss) self.assertIsInstance(obtained_prediction_box[0], Foo)
def _should_act_as_resource_variable(self): """Pass resource_variable_ops.is_resource_variable check.""" pass # Register a conversion function which reads the value of the variable, # allowing instances of the class to be used as tensors. def _tensor_conversion(var, dtype=None, name=None, as_ref=False): # Try to avoid assignments to and other mutations of MirroredVariable # state except through a DistributionStrategy.update() call. assert not as_ref return ops.internal_convert_to_tensor( var.get(), dtype=dtype, name=name, as_ref=as_ref) ops.register_tensor_conversion_function(DistributedVariable, _tensor_conversion) ops.register_dense_tensor_like_type(DistributedVariable) class _MirroredSaveable(saver.BaseSaverBuilder.ResourceVariableSaveable): """Class for defining how to restore a MirroredVariable.""" def __init__(self, mirrored_variable, primary_variable, name): self._mirrored_variable = mirrored_variable super(_MirroredSaveable, self).__init__(primary_variable, "", name) def restore(self, restored_tensors, restored_shapes): """Restore the same value into all variables.""" tensor, = restored_tensors return control_flow_ops.group([ _assign_on_device(d, v, tensor)
'of type {!r}'.format(dtype.name, self.dtype.name)) val = ops.internal_convert_to_tensor(self._variable, self._variable.dtype, name, as_ref=False) with ops.colocate_with(None, ignore_existing=True): with ops.device(val.device): return math_ops.cast(val, self.dtype) def _should_act_as_resource_variable(self): """Pass resource_variable_ops.is_resource_variable check.""" pass # TODO(reedwm): Define operator overloads. ops.register_tensor_conversion_function( AutoCastVariable, AutoCastVariable._dense_var_to_tensor) # pylint:disable=protected-access ops.register_dense_tensor_like_type(AutoCastVariable) # We have DistributedVariable subclass to pass # isinstance(..., DistributedVariable) checks when wrapping a # DistributedVariable. # TODO(reedwm): We should not wrap DistributedVariable, but instead have # DistributedVariable wrap AutoCastVariable. Subclassing DistributedVariable is # messy, because we do not fully implement the interface of DistributedVariable. class AutoCastDistributedVariable(AutoCastVariable, distribute_values.DistributedVariable): """Version of AutoCastVariable that subclasses DistributedVariable.""" def __init__(self, variable): if not isinstance(variable, distribute_values.DistributedValues):
def devices(self): return set(tensor.device for tensor in self.tensors) def __str__(self): return "PartitionedTensor([%s, ...], dtype=%s, shape=%s)" % ( self.tensors[0].name, self.dtype.name, tuple(self.shape.as_list())) def __hash__(self): return hash(tuple(self.tensors)) def as_tensor(self, dtype=None, name=None, as_ref=False): with ops.name_scope(name, "PartitionedTensor.as_tensor", self.tensors): assert not as_ref assert dtype in [None, self.dtype] result = array_ops.concat(self.tensors, axis=0) # Cache 'result' if we haven't already cached a value for this device. if result.device not in self._concats: self._concats[result.device] = result return self._concats[result.device] ops.register_tensor_conversion_function( PartitionedTensor, lambda val, dtype, name, as_ref: val.as_tensor(dtype, name, as_ref)) # TODO(b/69623235): Add a function for finding tensors that share gradients # to eliminate redundant fisher factor computations.
# Return an empty tensor so we only need to check for returned tensor # size being 0 as an indication of model ready. return array_ops.constant([], dtype=dtypes.string) else: # Get a 1-D boolean tensor listing whether each variable is initialized. variables_mask = math_ops.logical_not( array_ops.pack([state_ops.is_variable_initialized(v) for v in var_list]) ) # Get a 1-D string tensor containing all the variable names. variable_names_tensor = array_ops.constant([s.op.name for s in var_list]) # Return a 1-D tensor containing all the names of uninitialized variables. return array_ops.boolean_mask(variable_names_tensor, variables_mask) # pylint: disable=protected-access ops.register_tensor_conversion_function(Variable, Variable._TensorConversionFunction) Variable._OverloadAllOperators() ops.register_tensor_conversion_function(PartitionedVariable, PartitionedVariable._TensorConversionFunction) # pylint: enable=protected-access ops.register_dense_tensor_like_type(Variable) ops.register_proto_function( ops.GraphKeys.GLOBAL_VARIABLES, proto_type=variable_pb2.VariableDef, to_proto=Variable.to_proto, from_proto=Variable.from_proto, ) ops.register_proto_function( ops.GraphKeys.TRAINABLE_VARIABLES, proto_type=variable_pb2.VariableDef,
# Restore the original kernel_shape kernel_shape[2] = numKernelInputChannels for layerIndex in xrange(numStackedLayers - 1): if use_scope: commonScope = scope_name + '_shortcut_layer_' + str(layerIndex) blockInput = BatchNormedConvLayer(blockInput, ewma_trainer, kernel_shape, [1, 1, 1, 1], kernelInit, biasInitVal, 'SAME', operator, scale_after_norm, epsilon, is_train_phase, use_scope, commonScope) Hx = blockInput output = Tx * Hx + Cx * Gx if gateType == MywayFFLayer.RESIDUAL_GATE: # apply nonlinear op here # See the TODO above output = operator(output) return output # pylint: disable=protected-access ops.register_tensor_conversion_function(Layer, Layer._TensorConversionFunction) Layer._OverloadAllOperators() # pylint: enable=protected-access
... return tf.nn.embedding_lookup(self.sharded_variable.variables, x) >>> >>> model = Model() >>> model.fn(1).numpy() 2.0 >>> tf.saved_model.save(model, export_dir='/tmp/saved_model', ... signatures=model.serve_fn) """ @property def _type_spec(self): return ShardedVariableSpec( *(resource_variable_ops.VariableSpec(v.shape, v.dtype) for v in self._variables)) def _var_to_tensor(var, dtype=None, name=None, as_ref=False): del name if dtype is not None and not dtype.is_compatible_with(var.dtype): raise ValueError( 'Incompatible type conversion requested to type {!r} for variable ' 'of type {!r}'.format(dtype.name, var.dtype.name)) if as_ref: raise NotImplementedError( "ShardedVariable doesn't support being used as a reference.") return array_ops.concat(var.variables, axis=0) # Register a conversion function which reads the value of the variable, # allowing instances of the class to be used as tensors. ops.register_tensor_conversion_function(ShardedVariable, _var_to_tensor)
"""Pass resource_variable_ops.is_resource_variable check.""" pass def _dense_var_to_tensor(self, dtype=None, name=None, as_ref=False): """Converts a variable to a tensor.""" # pylint: disable=protected-access if _enclosing_tpu_context() is None: if hasattr(self._primary_var, '_dense_var_to_tensor'): return self._primary_var._dense_var_to_tensor(dtype, name, as_ref) else: return ops.convert_to_tensor(self._primary_var) # pylint: enable=protected-access if dtype is not None and dtype != self.dtype: return NotImplemented if as_ref: return self.handle else: return self.read_value() # Register a conversion function which reads the value of the variable, # allowing instances of the class to be used as tensors. def _tensor_conversion(var, dtype=None, name=None, as_ref=False): return var._dense_var_to_tensor(dtype=dtype, name=name, as_ref=as_ref) # pylint: disable=protected-access ops.register_tensor_conversion_function(ReplicatedVariable, _tensor_conversion) if not TF_23: ops.register_dense_tensor_like_type(ReplicatedVariable)
"""Base class for asset files which need to be tracked.""" def __init__(self, path): """Record the full path to the asset.""" # We use a variable here so that @tf.functions do not capture a literal # value. The init_scope prevents functions from capturing `path` in an # initialization graph, since it is transient and should not end up in a # serialized function body. When serialized in a SavedModel, the variable # will be set during the loading process to its location in the assets/ # directory. with ops.init_scope(): if context.executing_eagerly(): self._path = self._no_dependency( resource_variable_ops.ResourceVariable( path, dtype=dtypes.string, name="asset_path")) else: # Adding a variable is too disruptive when v1-style graph building, # since things may get fed and local variable initializers would then # need to be run. self._path = path @property def asset_path(self): """Fetch the current asset path.""" return self._path ops.register_tensor_conversion_function( TrackableAsset, lambda asset, **kw: ops.internal_convert_to_tensor(asset.asset_path, **kw))
"""Converts a variable to a tensor.""" with ds_context.enter_or_assert_strategy(self._distribute_strategy): return ops.convert_to_tensor(self._get(), dtype=dtype, name=name, as_ref=as_ref) # Register a conversion functions which reads the value of the variable, # allowing instances of the class to be used as tensors. # MirroredVariables def _tensor_conversion_mirrored(var, dtype=None, name=None, as_ref=False): return var._dense_var_to_tensor(dtype=dtype, name=name, as_ref=as_ref) # pylint: disable=protected-access ops.register_tensor_conversion_function(MirroredVariable, _tensor_conversion_mirrored) # Mirrored Values def _tensor_conversion_mirrored_val(value, dtype=None, name=None, as_ref=False): return ops.convert_to_tensor(value._get(), dtype=dtype, name=name, as_ref=as_ref) # pylint: disable=protected-access ops.register_tensor_conversion_function(Mirrored, _tensor_conversion_mirrored_val)
tc.register_type_abbreviation(ops.Tensor, 'tensorflow.Tensor') tc.register_type_abbreviation(dtypes.DType, 'tensorflow.DType') # core LabeledTensor types tc.register_type_abbreviation(Axis, 'labeled_tensor.Axis') tc.register_type_abbreviation(Axes, 'labeled_tensor.Axes') tc.register_type_abbreviation(LabeledTensor, 'labeled_tensor.LabeledTensor') @tc.returns(ops.Tensor) @tc.accepts(LabeledTensor) def _convert_labeled_tensor_to_tensor(value, *args, **kwargs): # call ops.convert_to_tensor to handle optional arguments appropriately return ops.convert_to_tensor(value.tensor, *args, **kwargs) ops.register_tensor_conversion_function( LabeledTensor, _convert_labeled_tensor_to_tensor) # tc class for anything that can be coerced into a LabeledTensor # pylint: disable=invalid-name LabeledTensorLike = tc.Union(LabeledTensor, ops.Tensor, np.ndarray, Scalar) # pylint: enable=invalid-name @tc.returns(LabeledTensor) @tc.accepts(LabeledTensorLike, object, tc.Optional(string_types)) def convert_to_labeled_tensor(value, dtype=None, name=None): """Converts the given `value` to a `LabeledTensor`. This function accepts `LabeledTensor` objects, 0-dimensional `Tensor` objects and numpy arrays, and Python scalars. Higher dimensional unlabeled tensors
raise NotImplementedError("surrogate_loss not implemented") @staticmethod def _tensor_conversion_function(v, dtype=None, name=None, as_ref=False): _ = name if dtype and not dtype.is_compatible_with(v.dtype): raise ValueError( "Incompatible type conversion requested to type '%s' for variable " "of type '%s'" % (dtype.name, v.dtype.name)) if as_ref: raise ValueError("%s: Ref type is not supported." % v) return v.value() # pylint: disable=protected-access ops.register_tensor_conversion_function( StochasticTensor, StochasticTensor._tensor_conversion_function) # pylint: enable=protected-access class _StochasticValueType(object): """Interface for the ValueType classes. This is the base class for MeanValue, SampleValue, SampleAndReshapeValue, and their descendants. """ def pushed_above(self, unused_value_type): pass def popped_above(self, unused_value_type): pass
def _read_variable_op(self): with ops.control_dependencies([self._parent_op]): return gen_resource_variable_ops.read_variable_op( self._handle, self._dtype) def set_shape(self, shape): self._shape = shape @property def op(self): """The op for this variable.""" return self._parent_op ops.register_tensor_conversion_function(_UnreadVariable, _dense_var_to_tensor) ops.register_dense_tensor_like_type(_UnreadVariable) # Register a conversion function which reads the value of the variable, # allowing instances of the class to be used as tensors. # Note: registering for Variable after ResourceVariable because inheritance will # otherwise lead to the wrong behavior. ops.register_tensor_conversion_function(ResourceVariable, _dense_var_to_tensor) ops.register_tensor_conversion_function( variables.Variable, variables.Variable._TensorConversionFunction) # pylint: disable=protected-access # pylint: disable=protected-access ResourceVariable._OverloadAllOperators() ops.register_dense_tensor_like_type(ResourceVariable)
return {checkpointable.VARIABLE_VALUE_KEY: _saveable_factory} # Register a conversion function which reads the value of the variable, # allowing instances of the class to be used as tensors. def _tensor_conversion_mirrored(var, dtype=None, name=None, as_ref=False): # Try to avoid assignments to and other mutations of MirroredVariable # state except through a DistributionStrategy.update() call. assert not as_ref return ops.internal_convert_to_tensor(var.get(), dtype=dtype, name=name, as_ref=as_ref) ops.register_tensor_conversion_function(MirroredVariable, _tensor_conversion_mirrored) class _TowerLocalSaveable(saver.BaseSaverBuilder.SaveableObject): """Class for defining how to restore a TowerLocalVariable.""" def __init__(self, tower_local_variable, name): self._tower_local_variable = tower_local_variable # We use a callable so that we don't have to evaluate this expression # in the case where we are trying to restore instead of save. def tensor(): return distribution_strategy_context.get_distribution_strategy( ).read_var(tower_local_variable) spec = saver.BaseSaverBuilder.SaveSpec( tensor=tensor,
if isinstance(tensor_or_op, ops.Tensor): op = tensor_or_op.op else: op = tensor_or_op return op.type == "Const" def _constant_tensor_conversion_function(v, dtype=None, name=None, as_ref=False): _ = as_ref return constant(v, dtype=dtype, name=name) ops.register_tensor_conversion_function( (list, tuple), _constant_tensor_conversion_function, 100) ops.register_tensor_conversion_function(np.ndarray, _constant_tensor_conversion_function, 100) ops.register_tensor_conversion_function(np.generic, _constant_tensor_conversion_function, 100) ops.register_tensor_conversion_function(object, _constant_tensor_conversion_function, 200) def _tensor_shape_tensor_conversion_function(s, dtype=None, name=None, as_ref=False):
captured_value = graph_placeholder(dtype=dtype or value.dtype, shape=value.shape, name=name) if captured_value.dtype == dtypes.resource: captured_value._handle_data = value._handle_data # pylint: disable=protected-access tensor_map[ops.tensor_id(value)] = (value, captured_value) else: captured_value = captured_value[1] return captured_value # TODO(apassos): it'd be really nice if we could scope this registration. # Note that we register this at a higher priority than ops.Tensor since we want # to handle subclass specific conversion before a superclass conversion. ops.register_tensor_conversion_function(tensor.Tensor, _convert_to_graph_constant, priority=-1) class _CapturingContext(object): """Tracks references to Tensors outside this context while it is active.""" def __init__(self): # known_ops are ops which are created while this context is active self.known_ops = set() # captured_tensors are all tensors referenced to by ops in this context but # not produced in it self.captured_tensors = set() def AddOp(self, op): # pylint: disable=invalid-name if op.type in ["Variable", "VariableV2", "VarHandleOp"]:
delta_get_op = delta_staging_area.get()[0] # Return the actual updates. The colocation constraint will be reapplied. return self.real_var.assign_sub(delta_get_op) @staticmethod # pylint: disable=bad-staticmethod-argument,invalid-name def _TensorConversionFunction(self, dtype=None, name=None, as_ref=False): """Utility function for converting a StagedModelVariable to a Tensor.""" del dtype, name # unused: this function returns the cached ref or value. if as_ref: return self._ref() else: return self._value() ops.register_tensor_conversion_function( StagedModelVariable, StagedModelVariable._TensorConversionFunction) # pylint: disable=protected-access class StagedVariableGetter(object): """A variable getter through staging buffers on devices. Instead of a caching device, this getter tracks where the variable is used. And on each device, it goes through a staging buffer. """ def __init__(self, device_num, devices, cpu_device, variable_mgr): """Initializer for StagedVariableGetter. Args: device_num: the current device index. devices: a list of all the devices to build towers.
tensor_value.tensor.CopyFrom( tensor_util.make_tensor_proto(value, dtype=dtype, shape=shape)) dtype_value = attr_value_pb2.AttrValue(type=tensor_value.tensor.dtype) const_tensor = g.create_op( "Const", [], [dtype_value.type], attrs={"value": tensor_value, "dtype": dtype_value}, name=name).outputs[0] return const_tensor @ops.RegisterShape("Const") def _ConstantShape(op): return [tensor_shape.TensorShape( [d.size for d in op.get_attr("value").tensor_shape.dim])] ops.register_tensor_conversion_function((list, tuple), constant, 100) ops.register_tensor_conversion_function(np.ndarray, constant, 100) ops.register_tensor_conversion_function(np.generic, constant, 100) ops.register_tensor_conversion_function(object, constant, 200) def _tensor_shape_tensor_conversion_function(s, dtype=None, name=None): if not s.is_fully_defined(): raise ValueError( "Cannot convert a partially known TensorShape to a Tensor: %s" % s) if dtype is not None: if dtype not in (types.int32, types.int64): raise TypeError("Cannot convert a TensorShape to dtype: %s" % dtype) else: dtype = types.int32 if name is None: name = "shape_as_tensor"
return self._var.scatter_min(sparse_delta, use_locking, name) def scatter_max(self, sparse_delta, use_locking=False, name=None): with ops.device(self._device): return self._var.scatter_max(sparse_delta, use_locking, name) def scatter_update(self, sparse_delta, use_locking=False, name=None): with ops.device(self._device): return self._var.scatter_update(sparse_delta, use_locking, name) def _dense_var_to_tensor(self, dtype=None, name=None, as_ref=False): with ops.device(self._device): return self._var._dense_var_to_tensor( # pylint: disable=protected-access dtype=dtype, name=name, as_ref=as_ref) def _as_graph_element(self): return self._var._as_graph_element() # pylint: disable=protected-access def _tensor_conversion_packed_var_and_device(var, dtype=None, name=None, as_ref=False): return var._dense_var_to_tensor(dtype=dtype, name=name, as_ref=as_ref) # pylint: disable=protected-access ops.register_tensor_conversion_function( PackedVarAndDevice, _tensor_conversion_packed_var_and_device)
tensor.Tensor(1, dtype=dtype)) def _lazy_zero_tensor(zero): return _zeros(zero.shape, zero.dtype) tensor.LazyZero.tensor = _lazy_zero_tensor def _lazy_zero_to_tensor(lazy_zero, dtype=None, name=None, as_ref=False): del as_ref, name, dtype return _zeros(lazy_zero.shape, lazy_zero.dtype) ops.register_tensor_conversion_function(tensor.LazyZero, _lazy_zero_to_tensor) def _indexed_slices_to_tensor(value): """Converts an IndexedSlices object `value` to a Tensor. Args: value: An ops.IndexedSlices object. Returns: A dense Tensor representing the values in the given IndexedSlices. Raises: ValueError: If the IndexedSlices does not have the same dtype. """ if value.dense_shape is None:
tensor.Tensor(shape, dtype=dtypes.int32), tensor.Tensor(1, dtype=dtype)) def _lazy_zero_tensor(zero): return _zeros(zero.shape, zero.dtype) tensor.LazyZero.tensor = _lazy_zero_tensor def _lazy_zero_to_tensor(lazy_zero, dtype=None, name=None, as_ref=False): del as_ref, name, dtype return _zeros(lazy_zero.shape, lazy_zero.dtype) ops.register_tensor_conversion_function(tensor.LazyZero, _lazy_zero_to_tensor) def _indexed_slices_to_tensor(value): """Converts an IndexedSlices object `value` to a Tensor. Args: value: An ops.IndexedSlices object. Returns: A dense Tensor representing the values in the given IndexedSlices. Raises: ValueError: If the IndexedSlices does not have the same dtype. """ if value.dense_shape is None:
# The created SavedModel is hermetic, it does not depend on # the original file and can be moved to another path. tf.io.gfile.remove("file.txt") tf.io.gfile.rename("/tmp/saved_model", "/tmp/new_location") reloaded_obj = tf.saved_model.load("/tmp/new_location") print(reloaded_obj.func()) ``` Attributes: asset_path: A 0-D `tf.string` tensor with path to the asset. """ def __init__(self, path): """Record the full path to the asset.""" # The init_scope prevents functions from capturing `path` in an # initialization graph, since it is transient and should not end up in a # serialized function body. with ops.init_scope(), ops.device("CPU"): self._path = ops.convert_to_tensor(path, dtype=dtypes.string, name="asset_path") @property def asset_path(self): """Fetch the current asset path.""" return self._path ops.register_tensor_conversion_function( Asset, lambda asset, **kw: ops.convert_to_tensor(asset.asset_path, **kw))
# pylint: disable=unused-argument,protected-access def _dense_var_to_tensor(var, dtype=None, name=None, as_ref=False): if dtype is not None and dtype != var.value().dtype: print("trying to switch the dtype to ", dtype, " from ", var.value().dtype) return NotImplemented return var.value() # pylint: enable=unused-argument,protected-access # Register a conversion function which reads the value of the variable, # allowing instances of the class to be used as tensors. ops.register_tensor_conversion_function(ResourceVariable, _dense_var_to_tensor) # pylint: disable=protected-access ResourceVariable._OverloadAllOperators() ops.register_dense_tensor_like_type(ResourceVariable) @ops.RegisterGradient("ReadVariableOp") def _ReadGrad(_, grad): """Gradient for read op.""" return grad @ops.RegisterGradient("ResourceGather") def _GatherGrad(op, grad): """Gradient for gather op."""