Tensorflow, Keras: How to create a trainable variable that only update in specific positions?

Question

For example, y=Ax

where A is an diagonal matrix, with its trainable weights (w1, w2, w3) on the diagonal.

A = [w1 ... ...
    ...  w2 ...
    ... ... w3]

How to create such trainable A in Tensorflow or Keras?

If I try A = tf.Variable(np.eye(3)), the total number of trainable weights would be 3*3=9, not 3. Because I only want to update (w1,w2,w3) that 3 weights.

A trick may be to use A = tf.Variable([1, 1, 1]) * np.eye(3), so that the 3 trainable weights are mapped into the diagonal of A.

My question is:

1. Would that trick work for my purpose? Would the gradient be correctly calculated?

2. What if the situation of A is more complicated? E.g. if I want to create:

where the w1, w2, ..., w6 are weights to be updated.

Answer 1

You have two different tools to address this problem.

1. You can create the variables you need and rearrange them into the desired form.
2. You can create more variables than you need then discard some to reach the desired form.

Both approach are not exclusive and you could you a mix of successives steps of type #1 and #2.

For example, for your first example (diagonal matrix), we can use approach #1.

w = tf.Variable(tf.zeros(n))
A = tf.diag(w) # creates a diagonal matrix with elements of w

For your second, more complex example, we could use approach #2.

A = tf.Variable(tf.zeros((n, n)))
A = tf.matrix_band_part(A, 1, 1) # keep only the central band of width 3
A = tf.matrix_set_diag(A, tf.ones(n)) # set diagonal to 1