def _make_transformed_lagrangian(lagrangian, activations_center, radius):
  """Returns a function that computes transformed Lagrangian L~(z).

  Args:
    lagrangian: function L(x), the lagrangian with fixed dual variables.
    activations_center: list of Device arrays corresponding to mean activations
      by layer.
    radius: list of Device arrays corresponding to the interval bound radii by
      layer.

  Returns:
    transformed_lagrangian: a function L~(z), defined as
      L~(z) = L(activations_center + z * radius).

  For compatibility with the paper, the inner optimization should have [-1, 1]
  element-wise constraints, so here we re-express:
    max_{x: a-rad<=x<=a+rad} L(x), as
    max_{z: -1<=z<=1} L~(z), with L~ defined as above.
  """
  def transformed_lagrangian(zs):
    zs = [a + z * r for (a, z, r) in zip(activations_center, zs, radius)]
    return lagrangian(zs)

  return transformed_lagrangian