`tf.svd` fails during GradientTape

user_3910261 0
Viewed 116

I'm trying to contract a network with multiple tensors and using singular value decomposition during contraction to simplify the contraction process. Whilst this works perfectly when I'm not taking any gradient, it fails once gradient tape starts to watch the tensors (I'm not sure why this is related). Below I wrote my simple contraction function and the function that I'm taking svd:

import tensorflow as tf

@tf.function
def contraction_step(network, max_singular_values: int):

    bottom = network[-1]
    uppper = network[-2]

    def contract_up_down(up,dn):
        shu = tf.shape(up)
        shd = tf.shape(dn)
        c = tf.einsum("ijkxlm,nkpyqr->injpxylqmr", up, dn)
        return tf.reshape(c, (
            shu[0]*shd[0], shu[1], shd[2], shu[-3], shd[-3], shu[-2]*shd[-2], shu[-1]*shd[-1]
        ))

    new = []
    multiplier = tf.eye(tf.shape(bottom[-1])[-1]*tf.shape(uppper[-1])[-1], dtype=uppper[-1].dtype)
    for ix in reversed(range(len(bottom))):
        tensor = contract_up_down(uppper[ix], bottom[ix])
        t = tf.einsum("ludpxor,ij->ludpxoj",tensor, multiplier)
        u, s, vh = svd(t, 1, max_singular_values = max_singular_values)
        multiplier = tf.tensordot(u, s/tf.norm(s), axes=(-1,0))
        new.insert(0, vh)
    new[-1] = tf.tensordot(new[-1], multiplier, axes=(-1,0))

    return network[:-2] + [new]

def svd(tensor,pivot,max_singular_values = None,cutoff = 0.0):
    left_dims = tf.shape(tensor)[:pivot]
    right_dims = tf.shape(tensor)[pivot:]
    tensor = tf.reshape(tensor, (tf.reduce_prod(left_dims), tf.reduce_prod(right_dims)))

    s, u, v = tf.linalg.svd(tensor)

    s_shape = tf.math.count_nonzero(
        tf.cast(s >= cutoff, dtype = tf.int32), dtype = tf.int32
    )
    if max_singular_values is None:
        max_singular_values = s_shape
    else:
        max_singular_values = tf.cast(tf.constant(max_singular_values), dtype = tf.int32)
    num_sing_vals_keep = tf.maximum(
        tf.minimum(max_singular_values, s_shape), tf.constant(1, dtype = tf.int32)
    )

    s = tf.slice(s, [0], [num_sing_vals_keep])
    u = tf.slice(u, [0, 0], [tf.shape(u)[0], num_sing_vals_keep])
    v = tf.slice(v, [0, 0], [tf.shape(v)[0], num_sing_vals_keep])

    vh = tf.linalg.adjoint(v)

    dim_s = tf.shape(s)[0]  # must use tf.shape (not s.shape) to compile
    u = tf.reshape(u, tf.concat([left_dims, [dim_s]], axis = -1))
    vh = tf.reshape(vh, tf.concat([[dim_s], right_dims], axis = -1))

    return u, tf.linalg.diag(s), vh

These functions work perfectly while using standalone:

upper = [tf.random.uniform((5,3,3,2,1,5), dtype=tf.float64) for _ in range(5)]
lower = [tf.random.uniform((5,3,3,2,1,5), dtype=tf.float64) for _ in range(5)]

contracted = contraction_step([upper, lower], 2)[0]
print(f"shapes: {', '.join([str(x.shape) for x in contracted])}")

# shapes: (2, 3, 3, 2, 2, 1, 2), (2, 3, 3, 2, 2, 1, 2), (2, 3, 3, 2, 2, 1, 2), (2, 3, 3, 2, 2, 1, 2), (2, 3, 3, 2, 2, 1, 2)

However, with the gradient, I get the following error:

with tf.GradientTape() as tape:
    tape.watch(upper + lower)
    contracted = contraction_step([upper, lower], 2)[0]

NotImplementedError: SVD gradient has not been implemented for input with unknown inner matrix shape.

It seems like for some reason during gradient mode TensorFlow loses the shape information of the tensors. Note that I get the same error when I set the tensors as tf.Variable instead of watching them manually. Any help would be highly appreciated!

Thanks

System information

  • OS Platform and Distribution: macOS v12.0.1
  • TensorFlow version: v2.6.0-rc2-32-g919f693420e 2.6.0
  • Python version: 3.8.9
svd gradient-descent tensorflow gradient python
1 Answers

I found a temporary solution that does not include all the aspects of the previous svd function but it works. TensorFlow requires object shapes to be set after slicing or manipulation (this might not be for every case but specific to mine). Thus I modified the svd function accordingly;

from functools import reduce
def svd(tensor, pivot, max_singular_values = None):

    multip = lambda x, y: x * y
    left_dims  = tensor.get_shape()[:pivot]#tf.shape(tensor)[:pivot]
    right_dims = tensor.get_shape()[pivot:]#tf.shape(tensor)[pivot:]
    tensor = tf.reshape(tensor, (reduce(multip, left_dims), reduce(multip, right_dims)))

    s, u, v = tf.linalg.svd(tensor)

    if max_singular_values is None:
        max_singular_values = s.shape[0]
    num_sing_vals_keep = min(s.shape[0], max_singular_values)
    
    s = tf.slice(s, [0], [num_sing_vals_keep])
    tf.ensure_shape(s, tf.TensorShape(num_sing_vals_keep))
    u = tf.slice(u, [0, 0], [tf.shape(u)[0], num_sing_vals_keep])
    v = tf.slice(v, [0, 0], [tf.shape(v)[0], num_sing_vals_keep])

    vh = tf.linalg.adjoint(v)

    dim_s = s.shape[0]#tf.shape(s)[0]
    u  = tf.reshape(u,  left_dims+(dim_s,))
    vh = tf.reshape(vh, (dim_s,) + right_dims)

    tf.ensure_shape(u, tf.TensorShape(left_dims+(dim_s,)))
    tf.ensure_shape(vh, tf.TensorShape((dim_s,) + right_dims))
    tf.ensure_shape(s, tf.TensorShape((dim_s,)))

    return u, tf.linalg.diag(s), vh

And this seems to be working both with tf.GradientTape() but I'm getting the following warnings when I run it in tf.vectorized_map

u,s,v = tf.vectorized_map(lambda vec: svd(vec, 1, 10), tf.random.uniform((10, 5, 5)))
WARNING:tensorflow:Using a while_loop for converting Svd

so if anyone has a better solution it's highly appreciated.

user_39102610
Related