def _run_joint_optimization(self, initial_states, inputs):
'''Finds multiple fixed points via a joint optimization over multiple
state vectors.
Args:
initial_states: Either an [n x n_states] numpy array or an
LSTMStateTuple with initial_states.c and initial_states.h as
[n_inits x n_states] numpy arrays. These data specify the initial
states of the RNN, from which the optimization will search for
fixed points. The choice of type must be consistent with state
type of rnn_cell.
inputs: A [n x n_inputs] numpy array specifying a set of constant
inputs into the RNN.
Returns:
fps: A FixedPoints object containing the optimized fixed points
and associated metadata.
'''
self._print_if_verbose('\tFinding fixed points '
'via joint optimization.')
n, _ = tf_utils.safe_shape(initial_states)
x, F = self._grab_RNN(initial_states, inputs)
# A shape [n,] TF Tensor of objectives (one per initial state) to be
# combined in _run_optimization_loop.
q = 0.5 * tf.reduce_sum(tf.square(F - x), axis=1)
xstar, F_xstar, qstar, dq, n_iters = self._run_optimization_loop(
q, x, F)
fps = FixedPoints(
xstar=xstar,
x_init=tf_utils.maybe_convert_from_LSTMStateTuple(initial_states),
inputs=inputs,
F_xstar=F_xstar,
qstar=qstar,
dq=dq,
n_iters=n_iters,
tol_unique=self.tol_unique,
dtype=self.np_dtype)
return fps
def _run_sequential_optimizations(self,
initial_states,
inputs,
q_prior=None):
'''Finds fixed points sequentially, running an optimization from one
initial state at a time.
Args:
initial_states: Either an [n x n_states] numpy array or an
LSTMStateTuple with initial_states.c and initial_states.h as
[n_inits x n_states] numpy arrays. These data specify the initial
states of the RNN, from which the optimization will search for
fixed points. The choice of type must be consistent with state
type of rnn_cell.
inputs: An [n x n_inputs] numpy array specifying a set of constant
inputs into the RNN.
q_prior (optional): An [n,] numpy array containing q values from a
previous optimization round. Provide these if performing
additional optimization iterations on a subset of outlier
candidate fixed points. Default: None.
Returns:
fps: A FixedPoints object containing the optimized fixed points
and associated metadata.
'''
is_fresh_start = q_prior is None
if is_fresh_start:
self._print_if_verbose('\tFinding fixed points via '
'sequential optimizations...')
n_inits, n_states = tf_utils.safe_shape(initial_states)
n_inputs = inputs.shape[1]
# Allocate memory for storing results
fps = FixedPoints(do_alloc_nan=True,
n=n_inits,
n_states=n_states,
n_inputs=n_inputs,
dtype=self.np_dtype)
for init_idx in range(n_inits):
index = slice(init_idx, init_idx + 1)
initial_states_i = tf_utils.safe_index(initial_states, index)
inputs_i = inputs[index, :]
if is_fresh_start:
self._print_if_verbose('\n\tInitialization %d of %d:' %
(init_idx + 1, n_inits))
else:
self._print_if_verbose(
'\n\tOutlier %d of %d (q=%.2e):' %
(init_idx + 1, n_inits, q_prior[init_idx]))
fps[init_idx] = self._run_single_optimization(
initial_states_i, inputs_i)
return fps
def find_fixed_points(self, initial_states, inputs):
'''Finds RNN fixed points and the Jacobians at the fixed points.
Args:
initial_states: Either an [n x n_dims] numpy array or an
LSTMStateTuple with initial_states.c and initial_states.h as
[n x n_dims] numpy arrays. These data specify the initial
states of the RNN, from which the optimization will search for
fixed points. The choice of type must be consistent with state
type of rnn_cell.
inputs: Either a [1 x n_inputs] numpy array specifying a set of
constant inputs into the RNN to be used for all optimization
initializations, or an [n x n_inputs] numpy array specifying
potentially different inputs for each initialization.
Returns:
unique_fps: A FixedPoints object containing the set of unique
fixed points after optimizing from all initial_states. Two fixed
points are considered unique if all absolute elementwise
differences are less than tol_unique AND the corresponding inputs
are unqiue following the same criteria. See FixedPoints.py for
additional detail.
all_fps: A FixedPoints object containing the likely redundant set
of fixed points (and associated metadata) resulting from ALL
initializations in initial_states (i.e., the full set of fixed
points before filtering out putative duplicates to yield
unique_fps).
'''
n = tf_utils.safe_shape(initial_states)[0]
self._print_if_verbose('\nSearching for fixed points '
'from %d initial states.\n' % n)
if inputs.shape[0] == 1:
inputs_nxd = np.tile(inputs, [n, 1]) # safe, even if n == 1.
elif inputs.shape[0] == n:
inputs_nxd = inputs
else:
raise ValueError('Incompatible inputs shape: %s.' % inputs.shape)
if self.method == 'sequential':
all_fps = self._run_sequential_optimizations(
initial_states, inputs_nxd)
elif self.method == 'joint':
all_fps = self._run_joint_optimization(initial_states, inputs_nxd)
else:
raise ValueError('Unsupported optimization method. Must be either \
\'joint\' or \'sequential\', but was \'%s\'' % self.method)
# Filter out duplicates after from the first optimization round
unique_fps = all_fps.get_unique()
self._print_if_verbose('\tIdentified %d unique fixed points.' %
unique_fps.n)
# Optionally run additional optimization iterations on identified
# fixed points with q values on the large side of the q-distribution.
if self.do_rerun_outliers:
unique_fps = self._run_additional_iterations_on_outliers(
unique_fps)
# Filter out duplicates after from the second optimization round
unique_fps = unique_fps.get_unique()
if self.do_compute_jacobians:
self._print_if_verbose('\tComputing Jacobian at %d '
'unique fixed points.' % unique_fps.n)
J_xstar = self._compute_multiple_jacobians_np(unique_fps)
unique_fps.J_xstar = J_xstar
self._print_if_verbose('\n\tFixed point finding complete.\n')
return unique_fps, all_fps