# Copyright OTT-JAX
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Any, Callable
import jax
import jax.numpy as jnp
import numpy as np
__all__ = ["fixpoint_iter", "fixpoint_iter_backprop"]
[docs]
def fixpoint_iter(
cond_fn: Callable[[int, Any, Any], bool],
body_fn: Callable[[Any, Any, Any, Any], Any], min_iterations: int,
max_iterations: int, inner_iterations: int, constants: Any, state: Any
):
"""Implementation of a fixed point loop.
This fixed point loop iterator applies ``body_fn`` to a tuple
``(iteration, constants, state, compute_error)`` to output a new state, using
context provided in iteration and constants.
``body_fn`` is iterated (inner_iterations -1) times, and one last time with
the ``compute_error`` flag to ``True``, indicating that additional
computational effort can be spent on recalculating the latest error
(``errors`` are stored as the first element of the state tuple).
upon termination of these ``inner_iterations``, the loop is continued if
iteration is smaller than ``min_iterations``, stopped if equal/larger than
``max_iterations``, and interrupted if ``cond_fn`` returns False.
Args:
cond_fn : termination condition function
body_fn : body loop instructions
min_iterations : lower bound on the total amount of fixed point iterations
max_iterations : upper bound on the total amount of fixed point iterations
inner_iterations : number of iterations ``body_fn`` will be executed
successively before calling ``cond_fn``.
constants : constant (during loop) parameters passed on to body
state : state variable
Returns:
outputs state returned by ``body_fn`` upon termination.
""" # noqa: D401
# If number of minimal iterations matches maximal number, force a scan instead
# of a while loop.
force_scan = (min_iterations == max_iterations)
compute_error_flags = jnp.arange(inner_iterations) == inner_iterations - 1
def max_cond_fn(iteration_state):
iteration, state = iteration_state
return jnp.logical_and(
iteration < max_iterations,
jnp.logical_or(
iteration < min_iterations, cond_fn(iteration, constants, state)
)
)
def unrolled_body_fn(iteration_state):
def one_iteration(iteration_state, compute_error):
iteration, state = iteration_state
state = body_fn(iteration, constants, state, compute_error)
iteration += 1
return (iteration, state), None
iteration_state, _ = jax.lax.scan(
one_iteration, iteration_state, compute_error_flags
)
return (iteration_state, None) if force_scan else iteration_state
if force_scan:
(_, state), _ = jax.lax.scan(
lambda carry, x: unrolled_body_fn(carry), (0, state),
None,
length=max_iterations // inner_iterations
)
else:
_, state = jax.lax.while_loop(max_cond_fn, unrolled_body_fn, (0, state))
return state
def fixpoint_iter_fwd(
cond_fn, body_fn, min_iterations, max_iterations, inner_iterations,
constants, state
):
"""Forward iteration of fixed point iteration to handle backpropagation.
The main difference with fixpoint_iter is the checkpointing, in variable
states, of the state variables as they are recorded through iterations, every
inner_iterations. This sequence of states will be used in the backward loop.
Args:
cond_fn : termination condition function
body_fn : body loop instructions
min_iterations : lower bound on the total amount of fixed point iterations
max_iterations : upper bound on the total amount of fixed point iterations
inner_iterations : number of iterations body_fn will be executed
successively before calling cond_fn.
constants : constant (during loop) parameters passed on to body
state : state variable
Returns:
outputs state returned by body_fn upon termination.
"""
force_scan = min_iterations == max_iterations
compute_error_flags = jnp.arange(inner_iterations) == inner_iterations - 1
states = jax.tree_util.tree_map(
lambda x: jnp.zeros(
(max_iterations // inner_iterations + 1,) + jnp.shape(x),
dtype=jax.dtypes.result_type(x)
), state
)
def max_cond_fn(iteration_states_state):
iteration, _, state = iteration_states_state
return jnp.logical_and(
iteration < max_iterations,
jnp.logical_or(
iteration < min_iterations, cond_fn(iteration, constants, state)
)
)
def unrolled_body_fn(iteration_states_state):
iteration, states, state = iteration_states_state
states = jax.tree_util.tree_map(
lambda states, state: jax.lax.dynamic_update_index_in_dim(
states, state, iteration // inner_iterations, 0
), states, state
)
def one_iteration(iteration_state, compute_error):
iteration, state = iteration_state
state = body_fn(iteration, constants, state, compute_error)
iteration += 1
return (iteration, state), None
iteration_state, _ = jax.lax.scan(
one_iteration, (iteration, state), compute_error_flags
)
iteration, state = iteration_state
out = (iteration, states, state)
return (out, None) if force_scan else out
if force_scan:
(iteration, states, state), _ = jax.lax.scan(
lambda carry, x: unrolled_body_fn(carry), (0, states, state),
None,
length=max_iterations // inner_iterations
)
else:
iteration, states, state = jax.lax.while_loop(
max_cond_fn, unrolled_body_fn, (0, states, state)
)
return state, (constants, iteration, states)
def fixpoint_iter_bwd(
cond_fn, body_fn, min_iterations, max_iterations, inner_iterations, res, g
):
"""Backward iteration of fixed point iteration, using checkpointed states."""
del cond_fn
force_scan = (min_iterations == max_iterations)
constants, iteration, states = res
# The tree may contain some python floats
g_constants = jax.tree_util.tree_map(
lambda x: jnp.zeros_like(x, dtype=x.dtype)
if isinstance(x, (np.ndarray, jnp.ndarray)) else 0, constants
)
def bwd_cond_fn(iteration_g_gconst):
iteration, _, _ = iteration_g_gconst
return iteration >= 0
def unrolled_body_fn_no_errors(iteration, constants, state):
compute_error_flags = jnp.zeros((inner_iterations,), dtype=bool)
def one_iteration(iteration_state, compute_error):
iteration, state = iteration_state
state = body_fn(iteration, constants, state, compute_error)
iteration += 1
return (iteration, state), None
iteration_state, _ = jax.lax.scan(
one_iteration, (iteration, state), compute_error_flags
)
_, state = iteration_state
return state
def unrolled_body_fn(iteration_g_gconst):
iteration, g, g_constants = iteration_g_gconst
state = jax.tree_util.tree_map(
lambda x: x[iteration // inner_iterations], states
)
_, pullback = jax.vjp(
unrolled_body_fn_no_errors, iteration, constants, state
)
_, gi_constants, g_state = pullback(g)
g_constants = jax.tree_util.tree_map(
lambda x, y: x + y, g_constants, gi_constants
)
out = (iteration - inner_iterations, g_state, g_constants)
return (out, None) if force_scan else out
if force_scan:
(_, g_state, g_constants), _ = jax.lax.scan(
lambda carry, x: unrolled_body_fn(carry), (0, g, g_constants),
None,
length=max_iterations // inner_iterations
)
else:
_, g_state, g_constants = jax.lax.while_loop(
bwd_cond_fn, unrolled_body_fn,
(iteration - inner_iterations, g, g_constants)
)
return g_constants, g_state
# definition of backprop friendly variant of fixpoint_iter.
fixpoint_iter_backprop = jax.custom_vjp(
fixpoint_iter, nondiff_argnums=(0, 1, 2, 3, 4)
)
fixpoint_iter_backprop.defvjp(fixpoint_iter_fwd, fixpoint_iter_bwd)
