Autodiff implementation for gradient calculation

I have worked through some papers about the autodiff algorithm to implement it for myself (for learning purposes). I compared my algorithm in test cases to the output of tensorflow and their outputs did not match in most cases. Therefor i worked through the tutorial from this side and implemented it with tensorflow operations just for the matrix multiplication operation since that was one of the operations that did not work:

gradient of matmul and unbroadcast method:

def gradient_matmul(node, dx, adj):
    # dx is needed to know which of both parents should be derived
    a = node.parents[0]
    b = node.parents[1]
    # the operation was node.tensor = tf.matmul(a.tensor, b,tensor)
    if a == dx or b == dx:
        # result depends on which of the parents is the derivative
        mm = tf.matmul(adj, tf.transpose(b.tensor)) if a == dx else \
                tf.matmul(tf.transpose(a.tensor), adj)
        return mm
    else: 
        return None

def unbroadcast(adjoint, node):
    dim_a = len(adjoint.shape)
    dim_b = len(node.shape)
    if dim_a > dim_b:
        sum = tuple(range(dim_a - dim_b))
        res = tf.math.reduce_sum(adjoint, axis = sum)
        return res
    return adjoint

And finally the gradient calculation autodiff algorithm:

def gradient(y, dx):
    working = [y]
    adjoints = defaultdict(float)
    adjoints[y] = tf.ones(y.tensor.shape)
    while len(working) != 0:
        curr = working.pop(0)
        if curr == dx:
            return adjoints[curr]
        if curr.is_store:
            continue
        adj = adjoints[curr]
        for p in curr.parents:
            # for testing with matrix multiplication as only operation
            local_grad = gradient_matmul(curr, p, adj)
            adjoints[p] = unbroadcast(tf.add(adjoints[p], local_grad), p.tensor)
            if not p in working:
                working.append(p)

Yet it produces the same output as my initial implementation. I constructed a matrix multiplication test case:

x = tf.constant([[[1.0, 1.0], [2.0, 3.0]], [[4.0, 5.0], [6.0, 7.0]]])
y = tf.constant([[3.0, -7.0], [-1.0, 5.0]])
z = tf.constant([[[1, 1], [2.0, 2]], [[3, 3], [-1, -1]]])
w = tf.matmul(tf.matmul(x, y), z)

Where w should be derived for each of the variables. Tensorflow calculates the gradient:

[<tf.Tensor: shape=(2, 2, 2), dtype=float32, numpy=
array([[[-22.,  18.],
        [-22.,  18.]],

       [[ 32., -16.],
        [ 32., -16.]]], dtype=float32)>, <tf.Tensor: shape=(2, 2), dtype=float32, numpy=
array([[66., -8.],
       [80., -8.]], dtype=float32)>, <tf.Tensor: shape=(2, 2, 2), dtype=float32, numpy=
array([[[  5.,   5.],
        [ -1.,  -1.]],

       [[ 18.,  18.],
        [-10., -10.]]], dtype=float32)>]

My implementation calculates:

[[[-5.  7.]
  [-5.  7.]]

 [[-5.  7.]
  [-5.  7.]]]
[[33. 22.]
 [54. 36.]]
[[[ 9.  9.]
  [14. 14.]]

 [[-5. -5.]
  [-6. -6.]]]

Maybe the problem is the difference between numpys dot and tensorflows matmul? But then i don't know to fix the gradient or unbroadcast for the tensorflow method... Thanks for taking the time to look over my code! :)

I found the error, the gradient matmul should have been:

def gradient_matmul(node, dx, adj):
    a = node.parents[0]
    b = node.parents[1]
    if a == dx:
        return tf.matmul(adj, b.tensor, transpose_b=True)
    elif b == dx:
        return tf.matmul(a.tensor, adj, transpose_a=True)
    else:
        return None

Since i only want to transpose the last 2 dimensions