"""Holds partition info used by initializer functions.

def __init__(self, full_shape, var_offset):

"""Constructor.

if not isinstance(full_shape, collections_lib.Sequence) or isinstance(

full_shape, six.string_types):

raise TypeError(

"`full_shape` must be a sequence (like tuple or list) instead of " +

type(full_shape).__name__)

if not isinstance(var_offset, collections_lib.Sequence) or isinstance(

var_offset, six.string_types):

raise TypeError(

"`var_offset` must be a sequence (like tuple or list) instead of " +

type(var_offset).__name__)

if len(var_offset) != len(full_shape):

raise ValueError(

"Expected equal length, but `var_offset` is of length {} while "

"full_shape is of length {}.".format(

len(var_offset), len(full_shape)))

for i in xrange(len(full_shape)):

offset = var_offset[i]

shape = full_shape[i]

if offset < 0 or offset >= shape:

raise ValueError(

"Expected 0 <= offset < shape but found offset={}, shape={} for "

"var_offset={}, full_shape={}".format(offset, shape, var_offset,

full_shape))

self._full_shape = full_shape

self._var_offset = var_offset

@property

def full_shape(self):

return self._full_shape

@property

def var_offset(self):

return self._var_offset

def single_offset(self, shape):

"""Returns the offset when the variable is partitioned in at most one dim.

single_slice_dim = self.single_slice_dim(shape)

# If this variable is not being partitioned at all, single_slice_dim() could

# return None.

if single_slice_dim is None:

return 0

return self.var_offset[single_slice_dim]

def single_slice_dim(self, shape):

"""Returns the slice dim when the variable is partitioned only in one dim.

if not isinstance(shape, collections_lib.Sequence) or isinstance(

shape, six.string_types):

raise TypeError(

"`shape` must be a sequence (like tuple or list) instead of " +

type(shape).__name__)

if len(shape) != len(self.full_shape):

raise ValueError(

"Expected equal length, but received shape={} of length {} while "

"self.full_shape={} is of length {}.".format(shape, len(

shape), self.full_shape, len(self.full_shape)))

for i in xrange(len(shape)):

if self.var_offset[i] + shape[i] > self.full_shape[i]:

raise ValueError(

"With self.var_offset={}, a partition of shape={} would exceed "

"self.full_shape={} in dimension {}.".format(

self.var_offset, shape, self.full_shape, i))

slice_dim = None

for i in xrange(len(shape)):

if shape[i] == self.full_shape[i]:

continue

if slice_dim is not None:

raise ValueError(

"Cannot use single_slice_dim() with shape={} and "

"self.full_shape={} since slice dim could be either dimension {} "

"or {}.".format(shape, self.full_shape, i, slice_dim))

slice_dim = i

"""Variable store that carries a number of named Variables.

def __init__(self):

"""Create a variable store."""

self._vars = {} # A dictionary of the stored TensorFlow variables.

self._partitioned_vars = {} # A dict of the stored PartitionedVariables.

self.variable_scopes_count = {} # Count re-used variable scopes.

def open_variable_scope(self, scope_name):

if scope_name in self.variable_scopes_count:

self.variable_scopes_count[scope_name] += 1

else:

self.variable_scopes_count[scope_name] = 1

def close_variable_subscopes(self, scope_name):

for k in self.variable_scopes_count:

if not scope_name or k.startswith(scope_name + "/"):

self.variable_scopes_count[k] = 0

def variable_scope_count(self, scope_name):

return self.variable_scopes_count.get(scope_name, 0)

def get_variable(self, name, shape=None, dtype=dtypes.float32,

initializer=None, regularizer=None, reuse=None,

trainable=True, collections=None, caching_device=None,

partitioner=None, validate_shape=True, use_resource=None,

custom_getter=None):

"""Gets an existing variable with these parameters or create a new one.

if custom_getter is not None and not callable(custom_getter):

raise ValueError(

"Passed a custom_getter which is not callable: %s" % custom_getter)

# If a *_ref type is passed in an error would be triggered further down the

# stack. We prevent this using base_dtype to get a non-ref version of the

# type, before doing anything else. When _ref types are removed in favour of

# resources, this line can be removed.

try:

dtype = dtype.base_dtype

except AttributeError:

# .base_dtype not existing means that we will try and use the raw dtype

# which was passed in - this might be a NumPy type which is valid.

pass

# This is the main logic of get_variable. However, custom_getter

# may override this logic. So we save it as a callable and pass

# it to custom_getter.

# Note: the parameters of _true_getter, and their documentation, match

# *exactly* item-for-item with the docstring of this method.

def _true_getter(name, shape=None, dtype=dtypes.float32, # pylint: disable=missing-docstring

initializer=None, regularizer=None, reuse=None,

trainable=True, collections=None, caching_device=None,

partitioner=None, validate_shape=True, use_resource=None):

is_scalar = shape is not None and not shape

# Partitioned variable case

if partitioner is not None and not is_scalar:

if not callable(partitioner):

raise ValueError(

"Partitioner must be callable, but received: %s" % partitioner)

with ops.name_scope(None):

return self._get_partitioned_variable(name=name,

shape=shape,

dtype=dtype,

initializer=initializer,

regularizer=regularizer,

reuse=reuse,

trainable=trainable,

collections=collections,

caching_device=caching_device,

partitioner=partitioner,

validate_shape=validate_shape,

use_resource=use_resource)

# Special case for partitioned variable to allow reuse without having to

# specify partitioner.

if (reuse is True and partitioner is None

and name in self._partitioned_vars):

return self._get_partitioned_variable(name=name,

shape=shape,

dtype=dtype,

initializer=initializer,

regularizer=regularizer,

reuse=reuse,

trainable=trainable,

collections=collections,

caching_device=caching_device,

partitioner=None,

validate_shape=validate_shape,

use_resource=use_resource)

# Single variable case

if "%s/part_0" % name in self._vars:

raise ValueError(

"No partitioner was provided, but a partitioned version of the "

"variable was found: %s/part_0. Perhaps a variable of the same "

"name was already created with partitioning?" % name)

return self._get_single_variable(

name=name, shape=shape, dtype=dtype,

initializer=initializer, regularizer=regularizer, reuse=reuse,

trainable=trainable, collections=collections,

caching_device=caching_device, validate_shape=validate_shape,

use_resource=use_resource)

if custom_getter is not None:

return custom_getter(

getter=_true_getter, name=name, shape=shape, dtype=dtype,

initializer=initializer, regularizer=regularizer,

reuse=reuse, trainable=trainable, collections=collections,

caching_device=caching_device, partitioner=partitioner,

validate_shape=validate_shape, use_resource=use_resource)

else:

return _true_getter(

name, shape=shape, dtype=dtype,

initializer=initializer, regularizer=regularizer,

reuse=reuse, trainable=trainable, collections=collections,

caching_device=caching_device, partitioner=partitioner,

validate_shape=validate_shape, use_resource=use_resource)

def _get_partitioned_variable(

self, name, partitioner, shape=None, dtype=dtypes.float32,

initializer=None, regularizer=None, reuse=None,

trainable=True, collections=None, caching_device=None,

validate_shape=True, use_resource=None):

"""Gets or creates a sharded variable list with these parameters.

initializing_from_value = initializer is not None and isinstance(

initializer, ops.Tensor)

reuse_without_partition = reuse is True and partitioner is None

if name in self._vars:

raise ValueError(

"A partitioner was provided, but an unpartitioned version of the "

"variable was found: %s. Perhaps a variable of the same name was "

"already created without partitioning?" % name)

shape = tensor_shape.as_shape(shape)

if initializing_from_value:

shape = shape.merge_with(initializer.get_shape())

if not reuse_without_partition:

if not shape.is_fully_defined():

raise ValueError("Shape of a new partitioned variable (%s) must be "

"fully defined, but instead was %s." % (name, shape))

if shape.ndims < 1:

raise ValueError("A partitioned Variable must have rank at least 1, "

"shape: %s" % shape)

partitions = partitioner(shape=shape, dtype=dtype)

if not isinstance(partitions, collections_lib.Sequence):

raise ValueError("Partitioner must return a sequence, but saw: %s"

% partitions)

if len(partitions) != shape.ndims:

raise ValueError(

"Partitioner returned a partition list that does not match the "

"Variable's rank: %s vs. %s" % (partitions, shape))

if any([p < 1 for p in partitions]):

raise ValueError(

"Partitioner returned zero partitions for some axes: %s" %

partitions)

should_check = reuse is not None

if name in self._partitioned_vars:

if should_check and not reuse:

raise ValueError(

"Partitioned variable with name %s already exists. Did you mean to "

"set reuse=True in VarScope?"

% name)

existing_var = self._partitioned_vars[name]

if not shape.is_compatible_with(existing_var.get_shape()):

raise ValueError(

"Trying to reuse partitioned variable %s, but specified shape %s "

"and found shape %s."

% (name, shape, existing_var.get_shape()))

if not dtype.is_compatible_with(existing_var.dtype):

raise ValueError(

"Trying to reuse partitioned variable %s, but specified dtype %s "

"and found dtype %s."

% (name, dtype.name, existing_var.dtype.name))

# pylint: disable=protected-access

if (not reuse_without_partition and

existing_var._get_partitions() != partitions):

raise ValueError(

"Trying to reuse partitioned variable %s, but specified partitions "

"%s and found partitions %s." %

(name, partitions, existing_var._get_partitions()))

# pylint: enable=protected-access

return existing_var

if should_check and reuse:

raise ValueError("PartitionedVariable %s does not exist, or was not "

"created with tf.get_variable(). Did you mean to set "

"reuse=None in VarScope?" % name)

slice_dim, slice_shape = _compute_slice_dim_and_shape(

shape.as_list(), partitions)

vs = []

num_slices = partitions[slice_dim]

num_slices_with_excess = shape[slice_dim].value % num_slices

slice_offset = [0] * shape.ndims

if "%s/part_0" % name in self._vars:

if "%s/part_%d" % (name, num_slices - 1) not in self._vars:

raise ValueError(

"Partitioner returned a different partitioning than what was "

"already found. Partitioner returned %d shards, and shard "

"%s/part_0 was found, but %s/part_%d was not."

% (num_slices, name, name, num_slices - 1))

if "%s/part_%d" % (name, num_slices) in self._vars:

raise ValueError(

"Partitioner returned a different partitioning than what was "

"already found. Partitioner returned %d shards, and shard "

"%s/part_0 was found, but so was the extra shard %s/part_%d."

% (num_slices, name, name, num_slices))

for i in xrange(num_slices):

var_shape = slice_shape[:]

var_offset = slice_offset[:]

partition_info = _PartitionInfo(

full_shape=shape.as_list(), var_offset=var_offset)

if i < num_slices_with_excess:

var_shape[slice_dim] += 1

slice_offset[slice_dim] += var_shape[slice_dim]

var_full_name = "%s/part_%d" % (name, i)

with ops.name_scope(var_full_name + "/PartitionedInitializer"):

# Create the tensor to initialize the variable with default value.

if initializer is None:

init, initializing_from_value = self._get_default_initializer(

name=name, shape=shape, dtype=dtype)

if initializing_from_value:

init_shape = None

else:

init_shape = var_shape

elif callable(initializer):

init = initializer

init_shape = var_shape

elif isinstance(initializer, ops.Tensor):

init = array_ops.slice(initializer, var_offset, var_shape)

# Use the dtype of the given tensor.

dtype = init.dtype.base_dtype

init_shape = None

else:

init = ops.convert_to_tensor(initializer, dtype=dtype)

init = array_ops.slice(init, var_offset, var_shape)

init_shape = None

with ops.name_scope(None):

var = self._get_single_variable(

name=var_full_name,

shape=init_shape,

dtype=dtype,

initializer=init,

partition_info=partition_info,

regularizer=regularizer,

reuse=reuse,

trainable=trainable,

collections=collections,

caching_device=caching_device,

validate_shape=validate_shape,

use_resource=use_resource)

# pylint: disable=protected-access

var._set_save_slice_info(variables.Variable.SaveSliceInfo(

name, shape.as_list(), var_offset, var_shape))

vs.append(var)

# pylint: enable=protected-access

# pylint: disable=protected-access

partitioned_var = variables.PartitionedVariable(name=name,

shape=shape,

dtype=dtype,

variable_list=vs,

partitions=partitions)

# pylint: enable=protected-access

self._partitioned_vars[name] = partitioned_var

return partitioned_var

def _get_single_variable(self,

name,

shape=None,

dtype=dtypes.float32,

initializer=None,

regularizer=None,

partition_info=None,

reuse=None,

trainable=True,

collections=None,

caching_device=None,

validate_shape=True,

use_resource=None,):

"""Get or create a single Variable (e.g. a shard or entire variable).

# Set to true if initializer is a constant.

initializing_from_value = False

if initializer is not None and not callable(initializer):

initializing_from_value = True

if shape is not None and initializing_from_value:

raise ValueError("If initializer is a constant, do not specify shape.")

should_check = reuse is not None

dtype = dtypes.as_dtype(dtype)

shape = tensor_shape.as_shape(shape)

if name in self._vars:

# Here we handle the case when returning an existing variable.

if should_check and not reuse:

tb = self._vars[name].op.traceback[::-1]

# Throw away internal tf entries and only take a few lines.

tb = [x for x in tb if "tensorflow/python" not in x[0]][:3]

raise ValueError("Variable %s already exists, disallowed."

" Did you mean to set reuse=True in VarScope? "

"Originally defined at:\n\n%s" % (

name, "".join(traceback.format_list(tb))))

found_var = self._vars[name]

if not shape.is_compatible_with(found_var.get_shape()):

raise ValueError("Trying to share variable %s, but specified shape %s"

" and found shape %s." % (name, shape,

found_var.get_shape()))

if not dtype.is_compatible_with(found_var.dtype):

dtype_str = dtype.name

found_type_str = found_var.dtype.name

raise ValueError("Trying to share variable %s, but specified dtype %s"

" and found dtype %s." % (name, dtype_str,

found_type_str))

return found_var

# The code below handles only the case of creating a new variable.

if should_check and reuse:

raise ValueError("Variable %s does not exist, or was not created with "

"tf.get_variable(). Did you mean to set reuse=None in "

"VarScope?" % name)

if not shape.is_fully_defined() and not initializing_from_value:

raise ValueError("Shape of a new variable (%s) must be fully defined, "

"but instead was %s." % (name, shape))

# Create the tensor to initialize the variable with default value.

if initializer is None:

initializer, initializing_from_value = self._get_default_initializer(

name=name, shape=shape, dtype=dtype)

# Clear control dependencies while creating the initializer.

with ops.control_dependencies(None):

if initializing_from_value:

init_val = initializer

variable_dtype = None

else:

# Instantiate initializer if provided initializer is a type object.

if isinstance(initializer, type(init_ops.Initializer)):

initializer = initializer(dtype=dtype)

init_val = lambda: initializer( # pylint: disable=g-long-lambda

shape.as_list(), dtype=dtype, partition_info=partition_info)

variable_dtype = dtype.base_dtype

# Create the variable.

if use_resource is None:

# Set the default value if unspecified.

use_resource = False

if use_resource:

v = resource_variable_ops.ResourceVariable(

initial_value=init_val,

name=name,

trainable=trainable,

collections=collections,

caching_device=caching_device,

dtype=variable_dtype,

validate_shape=validate_shape)

else:

v = variables.Variable(

initial_value=init_val,

name=name,

trainable=trainable,

collections=collections,

caching_device=caching_device,

dtype=variable_dtype,

validate_shape=validate_shape)

self._vars[name] = v

logging.vlog(1, "Created variable %s with shape %s and init %s", v.name,

format(shape), initializer)

# Run the regularizer if requested and save the resulting loss.

if regularizer:

with ops.colocate_with(v.op):

with ops.name_scope(name + "/Regularizer/"):

loss = regularizer(v)

if loss is not None:

logging.vlog(1, "Applied regularizer to %s and added the result %s "

"to REGULARIZATION_LOSSES.", v.name, loss.name)

ops.add_to_collection(ops.GraphKeys.REGULARIZATION_LOSSES, loss)

return v

# Initialize variable when no initializer provided

def _get_default_initializer(self, name, shape=None, dtype=dtypes.float32):

"""Provide a default initializer and a corresponding value.

# If dtype is DT_FLOAT, provide a uniform unit scaling initializer

if dtype.is_floating:

initializer = init_ops.glorot_uniform_initializer()

initializing_from_value = False

# If dtype is DT_INT/DT_UINT, provide a default value `zero`

# If dtype is DT_BOOL, provide a default value `FALSE`

elif dtype.is_integer or dtype.is_unsigned or dtype.is_bool:

initializer = init_ops.zeros_initializer()(

shape=shape, dtype=dtype.base_dtype)

initializing_from_value = True

# NOTES:Do we need to support for handling DT_STRING and DT_COMPLEX here?

else:

raise ValueError("An initializer for variable %s of %s is required"

% (name, dtype.base_dtype))

"""Variable scope object to carry defaults to provide to `get_variable`.

def __init__(self,

reuse,

name="",

initializer=None,

regularizer=None,

caching_device=None,

partitioner=None,

custom_getter=None,

name_scope="",

dtype=dtypes.float32,

use_resource=None):

"""Creates a new VariableScope with the given properties."""

self._name = name

self._initializer = initializer

self._regularizer = regularizer

self._reuse = reuse

self._caching_device = caching_device

self._partitioner = partitioner

self._custom_getter = custom_getter

self._name_scope = name_scope

self._dtype = dtype

self._use_resource = use_resource

@property

def name(self):

return self._name

@property

def original_name_scope(self):

return self._name_scope

@property

def reuse(self):

return self._reuse

@property

def initializer(self):

return self._initializer

@property

def dtype(self):

return self._dtype

@property

def use_resource(self):

return self._use_resource

@property

def regularizer(self):

return self._regularizer

@property

def caching_device(self):

return self._caching_device

@property

def partitioner(self):

return self._partitioner

@property

def custom_getter(self):

return self._custom_getter

def reuse_variables(self):

"""Reuse variables in this scope."""

self._reuse = True

def set_initializer(self, initializer):

"""Set initializer for this scope."""

self._initializer = initializer

def set_dtype(self, dtype):

"""Set data type for this scope."""

self._dtype = dtype

def set_use_resource(self, use_resource):

"""Sets whether to use ResourceVariables for this scope."""

self._use_resource = use_resource

def set_regularizer(self, regularizer):

"""Set regularizer for this scope."""

self._regularizer = regularizer

def set_caching_device(self, caching_device):

"""Set caching_device for this scope."""

self._caching_device = caching_device

def set_partitioner(self, partitioner):

"""Set partitioner for this scope."""

self._partitioner = partitioner

def set_custom_getter(self, custom_getter):

"""Set custom getter for this scope."""

self._custom_getter = custom_getter

def get_collection(self, name):

"""Get this scope's variables."""

scope = self._name + "/" if self._name else ""

return ops.get_collection(name, scope)

def trainable_variables(self):

"""Get this scope's trainable variables."""

return self.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES)

def global_variables(self):

"""Get this scope's global variables."""

return self.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)

def get_variable(self,

var_store,

name,

shape=None,

dtype=None,

initializer=None,

regularizer=None,

reuse=None,

trainable=True,

collections=None,

caching_device=None,

partitioner=None,

validate_shape=True,

use_resource=None,

custom_getter=None,):

"""Gets an existing variable with this name or create a new one."""

if regularizer is None:

regularizer = self._regularizer

if caching_device is None:

caching_device = self._caching_device

if partitioner is None:

partitioner = self._partitioner

if custom_getter is None:

custom_getter = self._custom_getter

if reuse is None:

reuse = self._reuse

full_name = self.name + "/" + name if self.name else name

# Variable names only depend on variable_scope (full_name here),

# not name_scope, so we reset it below for the time of variable creation.

with ops.name_scope(None):

# Check that `initializer` dtype and `dtype` are consistent before

# replacing them with defaults.

if (dtype is not None and initializer is not None and

not callable(initializer)):

init_dtype = ops.convert_to_tensor(initializer).dtype.base_dtype

if init_dtype != dtype:

raise ValueError("Initializer type '%s' and explicit dtype '%s' "

"don't match." % (init_dtype, dtype))

if initializer is None:

initializer = self._initializer

if dtype is None:

dtype = self._dtype

if use_resource is None:

use_resource = self._use_resource

return var_store.get_variable(

full_name, shape=shape, dtype=dtype, initializer=initializer,

regularizer=regularizer, reuse=reuse, trainable=trainable,

collections=collections, caching_device=caching_device,

partitioner=partitioner, validate_shape=validate_shape,

use_resource=use_resource, custom_getter=custom_getter)

def _get_partitioned_variable(self,

var_store,

name,

shape=None,

dtype=None,

initializer=None,

regularizer=None,

trainable=True,

collections=None,

caching_device=None,

partitioner=None,

validate_shape=True,

use_resource=None):

"""Gets an existing variable with this name or create a new one."""

if initializer is None:

initializer = self._initializer

if regularizer is None:

regularizer = self._regularizer

if caching_device is None:

caching_device = self._caching_device

if partitioner is None:

partitioner = self._partitioner

if dtype is None:

dtype = self._dtype

if use_resource is None:

use_resource = self._use_resource

if self._custom_getter is not None:

raise ValueError(

"Private access to _get_partitioned_variable is not allowed when "

"a custom getter is set. Current custom getter: %s. "

"It is likely that you're using create_partitioned_variables. "

"If so, consider instead using get_variable with a non-empty "

"partitioner parameter instead." % self._custom_getter)

if partitioner is None:

raise ValueError("No partitioner was specified")