def lanczos_bidiag_step(i, ls): """Extends the Lanczos bidiagonalization ls by one step.""" u = read_colvec(ls.u, i) r = operator.apply_adjoint(u) # The shape inference doesn't work across cond, save and reapply the shape. r_shape = r.get_shape() r = tf.cond(i > 0, lambda: r - ls.beta.read(i - 1) * read_colvec(ls.v, i - 1), lambda: r) r.set_shape(r_shape) if orthogonalize: v, alpha = orthogonalize_(i - 1, ls.v, r) else: v, alpha = util.l2normalize(r) p = operator.apply(v) - alpha * u if orthogonalize: u, beta = orthogonalize_(i, ls.u, p) else: u, beta = util.l2normalize(p) return i + 1, update_state(ls, i, u, v, alpha, beta)
def lanczos_bidiag_step(i, ls): """Extends the Lanczos bidiagonalization ls by one step.""" u = read_colvec(ls.u, i) r = operator.apply_adjoint(u) # The shape inference doesn't work across cond, save and reapply the shape. r_shape = r.get_shape() r = control_flow_ops.cond( i > 0, lambda: r - ls.beta.read(i - 1) * read_colvec(ls.v, i - 1), lambda: r) r.set_shape(r_shape) if orthogonalize: v, alpha = orthogonalize_(i - 1, ls.v, r) else: v, alpha = util.l2normalize(r) p = operator.apply(v) - alpha * u if orthogonalize: u, beta = orthogonalize_(i, ls.u, p) else: u, beta = util.l2normalize(p) return i + 1, update_state(ls, i, u, v, alpha, beta)
def testL2Norm(self): with self.test_session(): x_np = np.array([[2], [-3.], [5.]]) x_norm_np = np.linalg.norm(x_np) x_normalized_np = x_np / x_norm_np x = constant_op.constant(x_np) l2norm = util.l2norm(x) l2norm_squared = util.l2norm_squared(x) x_normalized, x_norm = util.l2normalize(x) self.assertAllClose(l2norm.eval(), x_norm_np) self.assertAllClose(l2norm_squared.eval(), np.square(x_norm_np)) self.assertAllClose(x_norm.eval(), x_norm_np) self.assertAllClose(x_normalized.eval(), x_normalized_np)
def testL2Norm(self): with self.test_session(): x_np = np.array([[2], [-3.], [5.]]) x_norm_np = np.linalg.norm(x_np) x_normalized_np = x_np / x_norm_np x = constant_op.constant(x_np) l2norm = util.l2norm(x) l2norm_squared = util.l2norm_squared(x) x_normalized, x_norm = util.l2normalize(x) self.assertAllClose(l2norm.eval(), x_norm_np) self.assertAllClose(l2norm_squared.eval(), np.square(x_norm_np)) self.assertAllClose(x_norm.eval(), x_norm_np) self.assertAllClose(x_normalized.eval(), x_normalized_np)
def lanczos_bidiag(operator, k, orthogonalize=True, starting_vector=None, name="lanczos_bidiag"): """Computes a Lanczos bidiagonalization for a linear operator. Computes matrices `U` of shape `[m, k+1]`, `V` of shape `[n, k]` and lower bidiagonal matrix `B` of shape `[k+1, k]`, that satisfy the equations `A * V = U * B` and `A' * U[:, :-1] = V * B[:-1, :]'`. The columns of `U` are orthonormal and form a basis for the Krylov subspace `K(A*A', U[:,0])`. The columns of `V` are orthonormal and form a basis for the Krylov subspace `K(A'*A, A' U[:,0])`. Args: operator: An object representing a linear operator with attributes: - shape: Either a list of integers or a 1-D `Tensor` of type `int32` of length 2. `shape[0]` is the dimension on the domain of the operator, `shape[1]` is the dimension of the co-domain of the operator. On other words, if operator represents an M x N matrix A, `shape` must contain `[M, N]`. - dtype: The datatype of input to and output from `apply` and `apply_adjoint`. - apply: Callable object taking a vector `x` as input and returning a vector with the result of applying the operator to `x`, i.e. if `operator` represents matrix `A`, `apply` should return `A * x`. - apply_adjoint: Callable object taking a vector `x` as input and returning a vector with the result of applying the adjoint operator to `x`, i.e. if `operator` represents matrix `A`, `apply_adjoint` should return `conj(transpose(A)) * x`. k: An integer or a scalar Tensor of type `int32`. Determines the maximum number of steps to run. If an invariant subspace is found, the algorithm may terminate before `k` steps have been run. orthogonalize: If `True`, perform full orthogonalization. If `False` no orthogonalization is performed. starting_vector: If not null, must be a `Tensor` of shape `[n]`. name: A name scope for the operation. Returns: output: A namedtuple representing a Lanczos bidiagonalization of `operator` with attributes: u: A rank-2 `Tensor` of type `operator.dtype` and shape `[operator.shape[0], k_actual+1]`, where `k_actual` is the number of steps run. v: A rank-2 `Tensor` of type `operator.dtype` and shape `[operator.shape[1], k_actual]`, where `k_actual` is the number of steps run. alpha: A rank-1 `Tensor` of type `operator.dtype` and shape `[k]`. beta: A rank-1 `Tensor` of type `operator.dtype` and shape `[k]`. """ def tarray(size, dtype, name): return tf.TensorArray( dtype=dtype, size=size, tensor_array_name=name, clear_after_read=False) # Reads a row-vector at location i in tarray and returns it as a # column-vector. def read_colvec(tarray, i): return tf.expand_dims(tarray.read(i), -1) # Writes an column-vector as a row-vecor at location i in tarray. def write_colvec(tarray, colvec, i): return tarray.write(i, tf.squeeze(colvec)) # Ephemeral class holding Lanczos bidiagonalization state: # u = left Lanczos vectors # v = right Lanczos vectors # alpha = diagonal of B_k. # beta = subdiagonal of B_k. # Notice that we store the left and right Lanczos vectors as the _rows_ # of u and v. This is done because tensors are stored row-major and # TensorArray only supports packing along dimension 0. lanzcos_bidiag_state = collections.namedtuple("LanczosBidiagState", ["u", "v", "alpha", "beta"]) def update_state(old, i, u, v, alpha, beta): return lanzcos_bidiag_state( write_colvec(old.u, u, i + 1), write_colvec(old.v, v, i), old.alpha.write(i, alpha), old.beta.write(i, beta)) def gram_schmidt_step(j, basis, v): """Makes v orthogonal to the j'th vector in basis.""" v_shape = v.get_shape() basis_vec = read_colvec(basis, j) v -= tf.batch_matmul(basis_vec, v, adj_x=True) * basis_vec v.set_shape(v_shape) return j + 1, basis, v def orthogonalize_once(i, basis, v): j = tf.constant(0, dtype=tf.int32) _, _, v = tf.while_loop(lambda j, basis, v: j < i, gram_schmidt_step, [j, basis, v]) return util.l2normalize(v) # Iterated modified Gram-Schmidt orthogonalization adapted from PROPACK. # TODO(rmlarsen): This is possibly the slowest implementation of # iterated Gram-Schmidt orthogonalization since the abacus. Move to C++. def orthogonalize_(i, basis, v): v_norm = util.l2norm(v) v_new, v_new_norm = orthogonalize_once(i, basis, v) # If the norm decreases more than 1/sqrt(2), run a second # round of MGS. See proof in: # B. N. Parlett, ``The Symmetric Eigenvalue Problem'', # Prentice-Hall, Englewood Cliffs, NJ, 1980. pp. 105-109 return tf.cond(v_new_norm < 0.7071 * v_norm, lambda: orthogonalize_once(i, basis, v), lambda: (v_new, v_new_norm)) def stopping_criterion(i, _): # TODO(rmlarsen): Stop if an invariant subspace is detected. return i < k def lanczos_bidiag_step(i, ls): """Extends the Lanczos bidiagonalization ls by one step.""" u = read_colvec(ls.u, i) r = operator.apply_adjoint(u) # The shape inference doesn't work across cond, save and reapply the shape. r_shape = r.get_shape() r = tf.cond( i > 0, lambda: r - ls.beta.read(i - 1) * read_colvec(ls.v, i - 1), lambda: r) r.set_shape(r_shape) if orthogonalize: v, alpha = orthogonalize_(i - 1, ls.v, r) else: v, alpha = util.l2normalize(r) p = operator.apply(v) - alpha * u if orthogonalize: u, beta = orthogonalize_(i, ls.u, p) else: u, beta = util.l2normalize(p) return i + 1, update_state(ls, i, u, v, alpha, beta) with tf.name_scope(name): dtype = operator.dtype if starting_vector is None: starting_vector = tf.random_uniform( operator.shape[:1], -1, 1, dtype=dtype) u0, _ = util.l2normalize(starting_vector) ls = lanzcos_bidiag_state( u=write_colvec(tarray(k + 1, dtype, "u"), u0, 0), v=tarray(k, dtype, "v"), alpha=tarray(k, dtype, "alpha"), beta=tarray(k, dtype, "beta")) i = tf.constant(0, dtype=tf.int32) _, ls = tf.while_loop(stopping_criterion, lanczos_bidiag_step, [i, ls]) return lanzcos_bidiag_state( tf.matrix_transpose(ls.u.pack()), tf.matrix_transpose(ls.v.pack()), ls.alpha.pack(), ls.beta.pack())
def orthogonalize_once(i, basis, v): j = tf.constant(0, dtype=tf.int32) _, _, v = tf.while_loop(lambda j, basis, v: j < i, gram_schmidt_step, [j, basis, v]) return util.l2normalize(v)
def orthogonalize_once(i, basis, v): j = constant_op.constant(0, dtype=dtypes.int32) _, _, v = control_flow_ops.while_loop(lambda j, basis, v: j < i, gram_schmidt_step, [j, basis, v]) return util.l2normalize(v)
def orthogonalize_once(i, basis, v): j = constant_op.constant(0, dtype=dtypes.int32) _, _, v = control_flow_ops.while_loop(lambda j, basis, v: j < i, gram_schmidt_step, [j, basis, v]) return util.l2normalize(v)