def test_nest_constant(self, dtype):
ndtype = dtype_util.as_numpy_dtype(dtype)
input_structure = (np.ones(4, dtype=ndtype), (np.eye(3, dtype=ndtype),
np.zeros(4,
dtype=ndtype)))
ones_like_structure = rk_util.nest_constant(input_structure)
tf.nest.assert_same_structure(input_structure, ones_like_structure)
flat_ones_like_structure = tf.nest.flatten(ones_like_structure)
for component in flat_ones_like_structure:
self.assertAllClose(component, tf.ones(shape=component.shape))
def grad_fn(*dresults, **kwargs):
"""Adjoint sensitivity method to compute gradients."""
dresults = tf.nest.pack_sequence_as(results, dresults)
dstates = dresults.states
# The signature grad_fn(*dresults, variables=None) is not valid Python 2
# so use kwargs instead.
variables = kwargs.pop('variables', [])
assert not kwargs # This assert should never fail.
# TODO(b/138304303): Support complex types.
with tf.name_scope('{}Gradients'.format(self._name)):
get_dtype = lambda x: x.dtype
def error_if_complex(dtype):
if dtype.is_complex:
raise NotImplementedError(
'The adjoint sensitivity method does '
'not support complex dtypes.')
state_dtypes = tf.nest.map_structure(
get_dtype, initial_state)
tf.nest.map_structure(error_if_complex, state_dtypes)
common_state_dtype = dtype_util.common_dtype(initial_state)
real_dtype = dtype_util.real_dtype(common_state_dtype)
# We add initial_time to ensure that we know where to stop.
result_times = tf.concat(
[[tf.cast(initial_time, real_dtype)], results.times],
0)
num_result_times = tf.size(result_times)
# First two components correspond to reverse and adjoint states.
# the last component is adjoint state for variables.
terminal_augmented_state = tuple([
rk_util.nest_constant(initial_state, 0.0),
rk_util.nest_constant(initial_state, 0.0),
tuple(
rk_util.nest_constant(variable, 0.0)
for variable in variables)
])
# The XLA compiler does not compile code which slices/indexes using
# integer `Tensor`s. `TensorArray`s are used to get around this.
result_time_array = tf.TensorArray(
results.times.dtype,
clear_after_read=False,
size=num_result_times,
element_shape=[]).unstack(result_times)
# TensorArray shape should not include time dimension, hence shape[1:]
result_state_arrays = [
tf.TensorArray( # pylint: disable=g-complex-comprehension
dtype=component.dtype,
size=num_result_times - 1,
element_shape=component.shape[1:]).unstack(
component)
for component in tf.nest.flatten(results.states)
]
result_state_arrays = tf.nest.pack_sequence_as(
results.states, result_state_arrays)
dresult_state_arrays = [
tf.TensorArray( # pylint: disable=g-complex-comprehension
dtype=component.dtype,
size=num_result_times - 1,
element_shape=component.shape[1:]).unstack(
component)
for component in tf.nest.flatten(dstates)
]
dresult_state_arrays = tf.nest.pack_sequence_as(
results.states, dresult_state_arrays)
def augmented_ode_fn(backward_time, augmented_state):
"""Dynamics function for the augmented system.
Describes a differential equation that evolves the augmented state
backwards in time to compute gradients using the adjoint method.
Augmented state consists of 3 components `(state, adjoint_state,
vars)` all evaluated at time `backward_time`:
state: represents the solution of user provided `ode_fn`. The
structure coincides with the `initial_state`.
adjoint_state: represents the solution of adjoint sensitivity
differential equation as discussed below. Has the same structure
and shape as `state`.
vars: represent the solution of the adjoint equation for variable
gradients. Represented as a `Tuple(Tensor, ...)` with as many
tensors as there are `variables`.
Adjoint sensitivity equation describes the gradient of the solution
with respect to the value of the solution at previous time t. Its
dynamics are given by
d/dt[adj(t)] = -1 * adj(t) @ jacobian(ode_fn(t, z), z)
Which is computed as:
d/dt[adj(t)]_i = -1 * sum_j(adj(t)_j * d/dz_i[ode_fn(t, z)_j)]
d/dt[adj(t)]_i = -1 * d/dz_i[sum_j(no_grad_adj_j * ode_fn(t, z)_j)]
where in the last line we moved adj(t)_j under derivative by
removing gradient from it.
Adjoint equation for the gradient with respect to every
`tf.Variable` theta follows:
d/dt[grad_theta(t)] = -1 * adj(t) @ jacobian(ode_fn(t, z), theta)
= -1 * d/d theta_i[sum_j(no_grad_adj_j * ode_fn(t, z)_j)]
Args:
backward_time: Floating `Tensor` representing current time.
augmented_state: `Tuple(state, adjoint_state, variable_grads)`
Returns:
negative_derivatives: Structure of `Tensor`s equal to backwards
time derivative of the `state` componnent.
adjoint_ode: Structure of `Tensor`s equal to backwards time
derivative of the `adjoint_state` component.
adjoint_variables_ode: Structure of `Tensor`s equal to backwards
time derivative of the `vars` component.
"""
# The negative signs disappears after the change of variables.
# The ODE solver cannot handle the case initial_time > final_time
# and hence a change of variables backward_time = -time is used.
time = -backward_time
state, adjoint_state, _ = augmented_state
with tf.GradientTape() as tape:
tape.watch(variables)
tape.watch(state)
derivatives = ode_fn(time, state)
adjoint_no_grad = tf.nest.map_structure(
tf.stop_gradient, adjoint_state)
negative_derivatives = rk_util.weighted_sum(
[-1.0], [derivatives])
def dot_prod(tensor_a, tensor_b):
return tf.reduce_sum(tensor_a * tensor_b)
# See docstring for details.
adjoint_dot_derivatives = tf.nest.map_structure(
dot_prod, adjoint_no_grad, derivatives)
adjoint_dot_derivatives = tf.squeeze(
tf.add_n(
tf.nest.flatten(adjoint_dot_derivatives)))
adjoint_ode, adjoint_variables_ode = tape.gradient(
adjoint_dot_derivatives, (state, tuple(variables)),
unconnected_gradients=tf.UnconnectedGradients.ZERO)
return negative_derivatives, adjoint_ode, adjoint_variables_ode
def reverse_to_result_time(n, augmented_state, _):
"""Integrates the augmented system backwards in time."""
lower_bound_of_integration = result_time_array.read(n)
upper_bound_of_integration = result_time_array.read(n -
1)
_, adjoint_state, adjoint_variable_state = augmented_state
initial_state = _read_solution_components(
result_state_arrays, input_state_structure, n - 1)
initial_adjoint = _read_solution_components(
dresult_state_arrays, input_state_structure, n - 1)
initial_adjoint_state = rk_util.weighted_sum(
[1.0, 1.0], [adjoint_state, initial_adjoint])
initial_augmented_state = (initial_state,
initial_adjoint_state,
adjoint_variable_state)
# TODO(b/138304303): Allow the user to specify the Hessian of
# `ode_fn` so that we can get the Jacobian of the adjoint system.
# TODO(b/143624114): Support higher order derivatives.
augmented_results = self._solve(
ode_fn=augmented_ode_fn,
initial_time=-lower_bound_of_integration,
initial_state=initial_augmented_state,
solution_times=[-upper_bound_of_integration],
batch_ndims=batch_ndims)
# Results added an extra time dim of size 1, squeeze it.
select_result = lambda x: tf.squeeze(x, [0])
result_state = augmented_results.states
result_state = tf.nest.map_structure(
select_result, result_state)
status = augmented_results.diagnostics.status
return n - 1, result_state, status
_, augmented_state, _ = tf.while_loop(
lambda n, _, status: (n >= 1) & tf.equal(status, 0),
reverse_to_result_time,
(num_result_times - 1, terminal_augmented_state, 0),
back_prop=False)
_, adjoint_state, adjoint_variables = augmented_state
return adjoint_state, list(adjoint_variables)
def vjp_bwd(results_constants, dresults, variables=()):
"""Adjoint sensitivity method to compute gradients."""
results, constants = results_constants
adjoint_solver = self._make_adjoint_solver_fn()
dstates = dresults.states
# TODO(b/138304303): Support complex types.
with tf.name_scope('{}Gradients'.format(self._name)):
get_dtype = lambda x: x.dtype
def error_if_complex(dtype):
if dtype_util.is_complex(dtype):
raise NotImplementedError(
'The adjoint sensitivity method does '
'not support complex dtypes.')
state_dtypes = tf.nest.map_structure(get_dtype, initial_state)
tf.nest.map_structure(error_if_complex, state_dtypes)
common_state_dtype = dtype_util.common_dtype(initial_state)
real_dtype = dtype_util.real_dtype(common_state_dtype)
# We add initial_time to ensure that we know where to stop.
result_times = tf.concat(
[[tf.cast(initial_time, real_dtype)], results.times], 0)
num_result_times = tf.size(result_times)
# First two components correspond to reverse and adjoint states.
# the last two component is adjoint state for variables and constants.
terminal_augmented_state = tuple([
rk_util.nest_constant(initial_state, 0.0),
rk_util.nest_constant(initial_state, 0.0),
tuple(
rk_util.nest_constant(variable, 0.0)
for variable in variables),
rk_util.nest_constant(constants, 0.0),
])
# The XLA compiler does not compile code which slices/indexes using
# integer `Tensor`s. `TensorArray`s are used to get around this.
result_time_array = tf.TensorArray(
results.times.dtype,
clear_after_read=False,
size=num_result_times,
element_shape=[]).unstack(result_times)
# TensorArray shape should not include time dimension, hence shape[1:]
result_state_arrays = [
tf.TensorArray( # pylint: disable=g-complex-comprehension
dtype=component.dtype,
size=num_result_times - 1,
clear_after_read=False,
element_shape=component.shape[1:]).unstack(component)
for component in tf.nest.flatten(results.states)
]
result_state_arrays = tf.nest.pack_sequence_as(
results.states, result_state_arrays)
dresult_state_arrays = [
tf.TensorArray( # pylint: disable=g-complex-comprehension
dtype=component.dtype,
size=num_result_times - 1,
clear_after_read=False,
element_shape=component.shape[1:]).unstack(component)
for component in tf.nest.flatten(dstates)
]
dresult_state_arrays = tf.nest.pack_sequence_as(
results.states, dresult_state_arrays)
def augmented_ode_fn(backward_time, augmented_state):
"""Dynamics function for the augmented system.
Describes a differential equation that evolves the augmented state
backwards in time to compute gradients using the adjoint method.
Augmented state consists of 4 components `(state, adjoint_state,
vars, constants)` all evaluated at time `backward_time`:
state: represents the solution of user provided `ode_fn`. The
structure coincides with the `initial_state`.
adjoint_state: represents the solution of the adjoint sensitivity
differential equation as discussed below. Has the same structure
and shape as `state`.
variables: represent the solution of the adjoint equation for
variable gradients. Represented as a `Tuple(Tensor, ...)` with as
many tensors as there are `variables` variable outside this
function.
constants: represent the solution of the adjoint equation for
constant gradients. Has the same structure and shape as
`constants` variable outside this function.
The adjoint sensitivity equation describes the gradient of the
solution with respect to the value of the solution at a previous
time t. Its dynamics are given by
d/dt[adj(t)] = -1 * adj(t) @ jacobian(ode_fn(t, z), z)
Which is computed as:
d/dt[adj(t)]_i = -1 * sum_j(adj(t)_j * d/dz_i[ode_fn(t, z)_j)]
d/dt[adj(t)]_i = -1 * d/dz_i[sum_j(no_grad_adj_j * ode_fn(t, z)_j)]
where in the last line we moved adj(t)_j under derivative by
removing gradient from it.
Adjoint equation for the gradient with respect to every
`tf.Variable` and constant theta follows:
d/dt[grad_theta(t)] = -1 * adj(t) @ jacobian(ode_fn(t, z), theta)
= -1 * d/d theta_i[sum_j(no_grad_adj_j * ode_fn(t, z)_j)]
Args:
backward_time: Floating `Tensor` representing current time.
augmented_state: `Tuple(state, adjoint_state, variable_grads)`
Returns:
negative_derivatives: Structure of `Tensor`s equal to backwards
time derivative of the `state` componnent.
adjoint_ode: Structure of `Tensor`s equal to backwards time
derivative of the `adjoint_state` component.
adjoint_variables_ode: Structure of `Tensor`s equal to backwards
time derivative of the `vars` component.
adjoint_constants_ode: Structure of `Tensor`s equal to backwards
time derivative of the `constants` component.
"""
# The negative signs disappears after the change of variables.
# The ODE solver cannot handle the case initial_time > final_time
# and hence a change of variables backward_time = -time is used.
time = -backward_time
state, adjoint_state, _, _ = augmented_state
# TODO(b/152464477): Doesn't work reliably in TF1.
def grad_fn(state, variables, constants):
del variables # We compute these gradients via the GradientTape
# capturing them.
derivatives = ode_fn(time, state, **constants)
adjoint_no_grad = tf.nest.map_structure(
tf.stop_gradient, adjoint_state)
negative_derivatives = rk_util.weighted_sum(
[-1.0], [derivatives])
def dot_prod(tensor_a, tensor_b):
return tf.reduce_sum(tensor_a * tensor_b)
# See docstring for details.
adjoint_dot_derivatives = tf.nest.map_structure(
dot_prod, adjoint_no_grad, derivatives)
adjoint_dot_derivatives = tf.squeeze(
tf.add_n(tf.nest.flatten(adjoint_dot_derivatives)))
return adjoint_dot_derivatives, negative_derivatives
values = (state, tuple(variables), constants)
((_, negative_derivatives),
gradients) = tfp_gradient.value_and_gradient(
grad_fn, values, has_aux=True, use_gradient_tape=True)
(adjoint_ode, adjoint_variables_ode,
adjoint_constants_ode) = tf.nest.map_structure(
lambda v, g: tf.zeros_like(v)
if g is None else g, values, tuple(gradients))
return (negative_derivatives, adjoint_ode,
adjoint_variables_ode, adjoint_constants_ode)
def make_augmented_state(n, prev_augmented_state):
"""Constructs the augmented state for step `n`."""
(_, adjoint_state, adjoint_variable_state,
adjoint_constant_state) = prev_augmented_state
initial_state = _read_solution_components(
result_state_arrays,
input_state_structure,
n - 1,
)
initial_adjoint = _read_solution_components(
dresult_state_arrays,
input_state_structure,
n - 1,
)
initial_adjoint_state = rk_util.weighted_sum(
[1.0, 1.0], [adjoint_state, initial_adjoint])
augmented_state = (
initial_state,
initial_adjoint_state,
adjoint_variable_state,
adjoint_constant_state,
)
return augmented_state
def reverse_to_result_time(n, augmented_state,
solver_internal_state, _):
"""Integrates the augmented system backwards in time."""
lower_bound_of_integration = result_time_array.read(n)
upper_bound_of_integration = result_time_array.read(n - 1)
initial_augmented_state = make_augmented_state(
n, augmented_state)
# TODO(b/138304303): Allow the user to specify the Hessian of
# `ode_fn` so that we can get the Jacobian of the adjoint system.
# TODO(b/143624114): Support higher order derivatives.
solver_internal_state = (
adjoint_solver.
_adjust_solver_internal_state_for_state_jump( # pylint: disable=protected-access
ode_fn=augmented_ode_fn,
initial_time=-lower_bound_of_integration,
initial_state=initial_augmented_state,
previous_solver_internal_state=
solver_internal_state,
previous_state=augmented_state,
))
augmented_results = adjoint_solver.solve(
ode_fn=augmented_ode_fn,
initial_time=-lower_bound_of_integration,
initial_state=initial_augmented_state,
solution_times=[-upper_bound_of_integration],
batch_ndims=batch_ndims,
previous_solver_internal_state=solver_internal_state,
)
# Results added an extra time dim of size 1, squeeze it.
select_result = lambda x: tf.squeeze(x, [0])
result_state = augmented_results.states
result_state = tf.nest.map_structure(
select_result, result_state)
status = augmented_results.diagnostics.status
return (n - 1, result_state,
augmented_results.solver_internal_state, status)
initial_n = num_result_times - 1
solver_internal_state = adjoint_solver._initialize_solver_internal_state( # pylint: disable=protected-access
ode_fn=augmented_ode_fn,
initial_time=result_time_array.read(initial_n),
initial_state=make_augmented_state(
initial_n, terminal_augmented_state),
)
_, augmented_state, _, _ = tf.while_loop(
lambda n, _as, _sis, status:
(n >= 1) & tf.equal(status, 0),
reverse_to_result_time,
(initial_n, terminal_augmented_state,
solver_internal_state, 0),
back_prop=False,
)
(_, adjoint_state, adjoint_variables,
adjoint_constants) = augmented_state
if variables:
return (adjoint_state,
adjoint_constants), list(adjoint_variables)
else:
return adjoint_state, adjoint_constants
def _step(self, solver_state, diagnostics, max_ode_fn_evals, ode_fn, atol,
rtol, safety, ifactor, dfactor):
"""Take an adaptive Runge-Kutta step.
Args:
solver_state: `_DopriSolverInternalState` - solver internal state.
diagnostics: `_DopriDiagnostics` - info on the current `_solve` call.
max_ode_fn_evals: Integer `Tensor` specifying the maximum number of ode_fn
evaluations.
ode_fn: Callable(t, y) -> dy_dt.
atol: Absolute tolerance allowed, see `_solve` method for details.
rtol: Relative tolerance allowed, see `_solve` method for details.
safety: Safety factor, see `_solve` method for details.
ifactor: Maximum factor by which the step size can increase, see `_solve`
method for details.
dfactor: Minimum factor by which the step size can decrease, see `_solve`
method for details.
Returns:
solver_state: `_RungeKuttaSolverInternalState` holding new solver state.
Note that the step might not advance the time if the error tolerance
criterias were not met. In this case step_size is decreased.
diagnostics: `_DopriDiagnostics` holding diagnostic values after RK step.
"""
y0, f0, _, t0, dt, interp_coeff = solver_state
assertion_ops = []
# TODO(dkochkov) Profile performance impact of `control_dependencies` here.
with tf.name_scope('assertions'):
if self._max_num_steps is not None:
check_max_num_steps = tf.debugging.assert_less(
diagnostics.num_ode_fn_evaluations, max_ode_fn_evals,
'max_num_steps exceeded')
assertion_ops.append(check_max_num_steps)
if self._validate_args:
check_underflow = tf.debugging.assert_greater(
t0 + dt, t0, 'underflow in dt')
assert_finite = functools.partial(
tf.debugging.assert_all_finite,
message='non-finite values in solution')
check_numerics = tf.nest.map_structure(assert_finite, y0)
assertion_ops.append(check_underflow)
assertion_ops.append(check_numerics)
with tf.control_dependencies(assertion_ops):
y1, f1, y1_error, k = rk_util.runge_kutta_step(
ode_fn, y0, f0, t0, dt, _TABLEAU)
with tf.name_scope('error_ratio'):
# We use the same criteria for accepting step as in scipy.
abs_y0 = tf.nest.map_structure(tf.abs, y0)
abs_y1 = tf.nest.map_structure(tf.abs, y1)
max_y_vals = tf.nest.map_structure(tf.math.maximum, abs_y0, abs_y1)
ones_nest = rk_util.nest_constant(abs_y0)
error_tol = rk_util.weighted_sum([atol, rtol],
[ones_nest, max_y_vals])
def scale_errors(error_component, tol_scale):
abs_square_error = rk_util.abs_square(error_component)
abs_square_tol_scale = rk_util.abs_square(tol_scale)
return tf.divide(abs_square_error, abs_square_tol_scale)
scaled_errors = tf.nest.map_structure(scale_errors, y1_error,
error_tol)
error_ratio = rk_util.nest_rms_norm(scaled_errors)
accept_step = error_ratio <= 1
with tf.name_scope('update/state'):
y_next = rk_util.nest_where(accept_step, y1, y0)
f_next = rk_util.nest_where(accept_step, f1, f0)
t_next = tf.where(accept_step, t0 + dt, t0)
new_coefficients = rk_util.rk_fourth_order_interpolation_coefficients(
y0, y1, k, dt, _TABLEAU)
interp_coeff = rk_util.nest_where(accept_step, new_coefficients,
interp_coeff)
dt_next = util.next_step_size(dt, self.ORDER, error_ratio, safety,
dfactor, ifactor)
solver_state = _RungeKuttaSolverInternalState(
y_next, f_next, t0, t_next, dt_next, interp_coeff)
diagnostics = diagnostics._replace(
num_ode_fn_evaluations=diagnostics.num_ode_fn_evaluations +
self.ODE_FN_EVALS_PER_STEP)
return solver_state, diagnostics
