Matrix Powers in TensorFlow

Viewed 274

Given an integer k and a symmetric matrix A (as tf.Variable), how to compute the k-th power of A

tf.matmul(A, tf.matmul(A, tf.matmul(A, ...)... )

most efficiently in TensorFlow?

3 Answers

This could be one approach to tackle this problem.
i. Convert matrix A into numpy ndarray(let's say B)
ii. Compute k-th power of B using: np.linalg.matrix_power(B, k)
iii. Convert the result back into tf.Variable

Here is a working code for the above-mentioned approach

import tensorflow as tf
import numpy as np
k = 2
A = tf.Variable([[1, -3], [2, 5]])
B = A.numpy()
M = np.linalg.matrix_power(B, k)
power = tf.Variable(M)
print(power)

Using tf.while_loop should be quite efficient:

import tensorflow as tf

k = 3
A = tf.Variable([[1, -3], [2, 5]])
result = tf.Variable(A)

i = tf.constant(0)
c = lambda i: tf.less(i, k - 1)

def body(i):
  result.assign(tf.matmul(A, result))
  return [tf.add(i, 1)]

_ = tf.while_loop(c, body, [i])

print(result)
<tf.Variable 'Variable:0' shape=(2, 2) dtype=int32, numpy= array([[-41, -75], [ 50, 59]], dtype=int32)>

You can also do this with fractions. If you need it to stay in tensor form use the aproximate solution with the taylor seiries of (1+x)^n and integer powers with 1 being the identity matrix

Not really sure what you will do with it but you can probably find something cool.

Related