def apply_gradients(self, metric_grads_and_vars):
"""Makes optimization step.
Args:
metric_grads_and_vars: zip(metric, grads, vars), metric is
a list of real valued tensors of shape (..., q, p, 2, q, p, 2),
grad is a list of real valued tensors of shape (..., q, p, 2),
vars is a list of real valued tf Variables of shape
(..., q, p, 2)"""
for mgv in metric_grads_and_vars:
metric, grad, var = mgv
var_c = m.real_to_complex(var)
grad_c = m.real_to_complex(grad)
rgrad_c = self.manifold.egrad_to_rgrad(metric, var_c, grad_c)
lr = tf.cast(self._hyper["learning_rate"], dtype=var_c.dtype)
new_var = self.manifold.retraction(var_c, -lr * rgrad_c)
var.assign(m.complex_to_real(new_var))
def body(i, loss):
with tf.GradientTape() as tape:
qc = manifolds.real_to_complex(q)
loss = tf.math.real(tf.linalg.trace(
tf.linalg.adjoint(qc) @ h @ qc))
grad = tape.gradient(loss, q)
opt.apply_gradients(zip([grad], [q]))
return i + 1, loss
def _resource_apply_dense(self, grad, var):
# Complex version of grad
complex_grad = m.real_to_complex(grad)
# MERA like update
_, u, v = tf.linalg.svd(adj(complex_grad))
var.assign(m.convert.complex_to_real(-v @ adj(u)))
def _resource_apply_dense(self, grad, var):
# Complex version of grad and var
complex_var = m.real_to_complex(var)
complex_grad = m.real_to_complex(grad)
# learning rate
eps = tf.cast(self._get_hyper("eps"), dtype=complex_grad.dtype)
# noise
noise = tf.random.normal(var.shape, dtype=var.dtype)
noise = m.real_to_complex(noise)
noise = noise * math.sqrt(eps)
# search direction
r = -self.manifold.egrad_to_rgrad(complex_var,
0.5 * eps * complex_grad + noise)
# New value of var
new_var = self.manifold.retraction(complex_var, r)
# Update of var
var.assign(m.complex_to_real(new_var))
def _egrad_to_rgrad(self):
"""
Checking egrad_to_rgrad method.
1) rgrad is in the tangent space of a manifold's point
2) <v1 egrad> = <v1 rgrad>_m (matching between egrad and rgrad)
Args:
Returns:
list with two tf scalars that give maximum
violation of two conditions
"""
# vector that plays the role of a gradient
xi = tf.random.normal(self.u.shape + (2,))
xi = manifolds.real_to_complex(xi)
xi = tf.cast(xi, dtype=self.u.dtype)
# rgrad
rgrad = self.m.egrad_to_rgrad(self.u, xi)
err1 = rgrad - self.m.proj(self.u, rgrad)
if self.m.rank == 2:
err1 = tf.math.real(tf.linalg.norm(err1, axis=(-2, -1)))
elif self.m.rank == 3:
err1 = tf.math.real(rank3_norm(err1))
if self.m.rank == 2:
err2 = tf.reduce_sum(tf.math.conj(self.v1) * xi, axis=(-2, -1)) -\
self.m.inner(self.u, self.v1, rgrad)[..., 0, 0]
elif self.m.rank == 3:
err2 = tf.reduce_sum(tf.math.conj(self.v1) * xi, axis=(-3, -2, -1)) -\
self.m.inner(self.u, self.v1, rgrad)[..., 0, 0, 0]
err2 = tf.abs(tf.math.real(err2))
err1 = tf.reduce_max(err1)
err2 = tf.reduce_max(err2)
return err1, err2
def _egrad_to_rgrad(self):
"""
Checking egrad to rgrad.
1) rgrad is in a tangent space
2) <v1 egrad> = inner<v1 rgrad>
Args:
Returns:
error 1), dtype float32
error 2), dtype float32
"""
xi = tf.random.normal(self.u.shape + (2, ))
xi = manifolds.real_to_complex(xi)
xi = tf.cast(xi, dtype=self.u.dtype)
rgrad = self.m.egrad_to_rgrad(self.u, xi)
err1 = rgrad - self.m.proj(self.u, rgrad)
err1 = tf.abs(tf.linalg.norm(err1))
err2 = tf.reduce_sum(tf.math.conj(self.v1) * xi, axis=(-2, -1)) -\
self.m.inner(self.u, self.v1, rgrad)
err2 = tf.abs(tf.math.real(err2))
return tf.cast(err1, dtype=tf.float32), tf.cast(err2, dtype=tf.float32)
def _resource_apply_dense(self, grad, var):
self.iter = self.iter + 1
# Complex version of grad and var
complex_var = m.real_to_complex(var)
complex_grad = m.real_to_complex(grad)
# learning rate
lr = tf.cast(self._get_hyper("learning_rate"), complex_grad.dtype)
# Riemannian gradient
rgrad = self.manifold.egrad_to_rgrad(complex_var, complex_grad)
# Complex versions of m and v
momentum = self.get_slot(var, "momentum")
v = self.get_slot(var, "v")
if self.ams:
v_hat = self.get_slot(var, "v_hat")
v_hat_complex = m.real_to_complex(v_hat)
momentum_complex = m.real_to_complex(momentum)
v_complex = m.real_to_complex(v)
# Update m, v and v_hat
beta1 = tf.cast(self._get_hyper("beta1"), dtype=momentum_complex.dtype)
beta2 = tf.cast(self._get_hyper("beta2"), dtype=momentum_complex.dtype)
momentum_complex = beta1 * momentum_complex +\
(1 - beta1) * rgrad
v_complex = beta2 * v_complex +\
(1 - beta2) * self.manifold.inner(complex_var,
rgrad,
rgrad)
if self.ams:
v_hat_complex = tf.maximum(tf.math.real(v_complex),
tf.math.real(v_hat_complex))
v_hat_complex = tf.cast(v_hat_complex, dtype=v_complex.dtype)
# Bias correction
lr_corr = lr * tf.math.sqrt(1 - beta2 ** self.iter) /\
(1 - beta1 ** self.iter)
# New value of var
if self.ams:
# Search direction
search_dir = -lr_corr * momentum_complex /\
(tf.sqrt(v_hat_complex) + self.eps)
new_var, momentum_complex =\
self.manifold.retraction_transport(complex_var,
momentum_complex,
search_dir)
else:
# Search direction
search_dir = - lr_corr * momentum_complex /\
(tf.sqrt(v_complex) + self.eps)
new_var, momentum_complex =\
self.manifold.retraction_transport(complex_var,
momentum_complex,
search_dir)
# Assigning new value of momentum
momentum.assign(m.complex_to_real(momentum_complex))
# Assigning new value of v and v_hat
v.assign(m.complex_to_real(v_complex))
if self.ams:
v_hat.assign(m.complex_to_real(v_hat_complex))
# Update of var
var.assign(m.complex_to_real(new_var))
def _resource_apply_dense(self, grad, var):
# Complex version of grad and var
complex_var = m.real_to_complex(var)
complex_grad = m.real_to_complex(grad)
# learning rate
lr = tf.cast(self._get_hyper("learning_rate"),
dtype=complex_grad.dtype)
# Riemannian gradient
rgrad = self.manifold.egrad_to_rgrad(complex_var, complex_grad)
# Upadte of vars (step and retruction)
if self._momentum:
if self._use_nesterov:
# Update momentum
momentum_var = self.get_slot(var, "momentum")
momentum_complex_old = m.real_to_complex(momentum_var)
momentum = tf.cast(self._get_hyper("momentum"),
dtype=momentum_complex_old.dtype)
momentum_complex_new = momentum * momentum_complex_old +\
(1 - momentum) * rgrad
# Transport and retruction
new_var, momentum_complex =\
self.manifold.retraction_transport(complex_var,
momentum_complex_new,
-lr * momentum_complex_new -\
lr * momentum *\
(momentum_complex_new -\
momentum_complex_old))
momentum_var.assign(m.complex_to_real(momentum_complex))
else:
# Update momentum
momentum_var = self.get_slot(var, "momentum")
momentum_complex = m.real_to_complex(momentum_var)
momentum = tf.cast(self._get_hyper("momentum"),
dtype=momentum_complex.dtype)
momentum_complex = momentum * momentum_complex +\
(1 - momentum) * rgrad
# Transport and retruction
new_var, momentum_complex =\
self.manifold.retraction_transport(complex_var,
momentum_complex,
-lr * momentum_complex)
momentum_var.assign(m.complex_to_real(momentum_complex))
else:
# New value of var
new_var = self.manifold.retraction(complex_var, -lr * rgrad)
# Update of var
var.assign(m.complex_to_real(new_var))