I am trying to do Strassen’s matrix multiplication on TensorFlow. I am trying to perform a graph execution using the following code
import tensorflow as tf

@tf.function  # The decorator converts `split_tf` into a `Function`.
def split_tf(matrix):
  shape = tf.shape(matrix)
  shape2 = tf.cast(tf.divide(shape, 2), dtype=tf.int32)
  return tf.slice(A, [0, 0], shape2), tf.slice(A, [0, shape2[1]], shape2), 
         tf.slice(A, [shape2[0], 0], shape2), 
         tf.slice(A, [shape2[0], shape2[1]], shape2)

@tf.function  # The decorator converts `strassen_tf` into a `Function`.
def strassen_tf(x, y):
  # Base case when size of matrices is 1x1
  if tf.rank(x) == 1:
      return tf.math.multiply(x, y)

  # Splitting the matrices into quadrants. This will be done recursively
  # until the base case is reached.

  a, b, c, d = split_tf(x)
  e, f, g, h = split_tf(y)

  # Computing the 7 products, recursively (p1, p2...p7)
  p1 = strassen_tf(a, tf.math.subtract(f, h)) 
  p2 = strassen_tf(tf.math.add(a, b), h)
  p3 = strassen_tf(tf.math.add(c, d), e)
  p4 = strassen_tf(d, tf.math.subtract(g, e))
  p5 = strassen_tf(tf.math.add(a, d), tf.math.add(e, h))
  p6 = strassen_tf(tf.math.subtract(b, d), tf.math.add(g, h))
  p7 = strassen_tf(tf.math.subtract(a, c), tf.math.add(e, f))

  # Computing the values of the 4 quadrants of the final matrix c
  c11 = tf.math.add(tf.math.subtract(tf.math.add(p5, p4), p2), p6)
  c12 = tf.math.add(p1, p2)          
  c21 = tf.math.add(p3, p4)
  c22 = tf.math.subtract(tf.math.subtract(tf.math.add(p1, p5), p3), p7)

  # Combining the 4 quadrants into a single matrix by stacking horizontally and vertically.
  c = tf.concat([tf.concat([C11, C12], axis=1), tf.concat([C21, C22], axis=1)], axis=0)

  return c

Now when I use the above code to multiply two tensors as below:
n = 2;
A = tf.random.normal([n, n], 0, 1, dtype=tf.float64)
B = tf.random.normal([n, n], 0, 1, dtype=tf.float64)
D = strassen_tf(A, B)

the code stuck at the last statement.
Thanks in advance!