Python 在带有掩码的tensorflow 中执行矩阵乘法时出错

Let's say val is a matrix of size (2, N). I need to multiply it with a mask matrix mask of size
(K,K) containing values zeros and ones at different indexes. This should output a matrix result of size (N, K, K) matrix with each submatrix of result along the dimension 0 to be a (K,K) matrix where zeros are replaced by val(i,1) and ones are replaced by val(i,2).

例如,

mask = tf.constant([[0, 1, 0, 1],
                   [1, 0, 0, 1],
                   [1, 1, 1, 0],
                   [0, 1, 0, 0]], dtype=tf.int32)

val = tf.constant([[3, 2, 8, 1, 9, 5, 6], [7, 4, 9, 8, 3, 1, 9]])

那么输出应该是这样的7 x 4 x 4矩阵，

result  = 

tf.Tensor(
[[[ 3.  7.  3.  7.]
  [ 7.  3.  3.  7.]
  [ 7.  7.  7.  3.]
  [ 3.  7.  3.  3.]]

 [[ 2.  4.  2.  4.]
  [ 4.  2.  2.  4.]
  [ 4.  4.  4.  2.]
  [ 3.  4.  3.  2.]]
          :
          :
          :
 [[ 6.  9.  6.  9.]
  [ 9.  6.  6.  9.]
  [ 9.  9.  9.  6.]
  [ 6.  9.  6.  6.]]]

目前，我正在使用for循环迭代val的每一列，以执行以下操作

val[0,i]*mask + val[1,i]*(1-mask)个

我希望将其矢量化，以纳入TensorFlow的矩阵乘法能力.