Replacing tf.pad with vector operations

[pure knot]

I'm working on a model to produce board evaluations for chess. I'm inspired by GNNs/MPNNs and want to meaningfully represent the chess moves that can be performed on a board by their relationship with other squares. I embed all the squares on a board based on the piece occupying them, resulting in a board embedding of (64, piece_embedding_nfeatures). For each type of move possible (rook moves, knight moves, bishop moves), I create a tf.Variable translation_var of shape (piece_embedding_nfeatures, piece_embedding_nfeatures) and then for each pair of squares reachable by this move type on the board I perform matrix multiplication the source square embedding and update the board state on the destination square. My hope is that upon recursion, the network will be able to represent the relationship between peices and squares (this queen is under attack, this square is defended by this pawn, etc.). My current solution for implementing this is to apply tf.pad operations to each translation_var for all the source and destionation squares possible, and than sum the result to produce a (64 * piece_embedding_nfeatures, 64 * piece_embedding_nfeatures) which is used in matrix multiplication on the board state. However, the repeated tf.pad solution is enormously inefficent both to construct and evauluate the graph. I'm hoping that there's a solution with vector operations that can be more efficient.

Here it is my current implementation:

class Translation(keras.layers.Layer):
   ...

    @staticmethod
    def get_translations(
        units: int, translation_vars: Iterable[tuple[tf.Tensor, Iterable[tf.Tensor]]]
    ):
        """Gets translations matrix"""

        var_translations = [tf.eye(64 * units)]
        for var, edge_paddings in translation_vars:
            for paddings in edge_paddings:
                var_translations.append(tf.pad(var, paddings, "CONSTANT"))

        return tf.add_n(var_translations)

Here units is a positive integer equal to the number of piece/square embedding features. For each element of translation_vars, the first value is a tf.Variable of shape (units, units), and the second value is a list of tf.pad arguments which translate the var matrix to the appropriate source and dest offsets. Upon summing all the padded vars, we produce the final matrix that is matmuled against the board state.

[pure knot]

I feel like there is a solution with tf.tensordot outer product, but I can't quite wrap my head around some of the details. If I do something like this:

piece_embedding_nfeatures = 8
squares = 64

source_square = 1
dest_square = 0
var = tf.eye(piece_embedding_nfeatures)

edge = tf.tensor_scatter_nd_update(
  tf.zeros([squares, squares]),
  [[source_square, dest_square]],
  [1.]
)
m = tf.tensordot(var, edge, 0)

m "feels" so close to what I want, but the shape is (piece_embedding_nfeatures, piece_embedding_nfeatures, squares, squares) . It has the same number of elements as my desired matix. But I'm struggling to figure out how to reshape correctly. In this example, using tf.eye I would like it to result in a matrix such that tf.matmul(board, m) would "move" the piece from the source_square to the dest_square. But if I just tf.reshape(m, (piece_embedding_nfeatures * squares, piece_embedding_nfeatures * squares)) it doesn't work at all and the 1. s from var are mapped all over the place