tensorflowで各バッチ要素内計算のやり方

Question

tensorflowのグラフ内で以下の計算を行いたいです。

a はテンソル shape: [2, 2, 2]
a_mean は a の0階成分の平均 shape:[2]

a = tf.constant(
        [[[1.0, 2.0], 
          [3.0, 4.0]], 

         [[1.0, 3.0], 
         [[5.0, 7.0]]]
)

a_mean = tf.reduce_mean(a, [1, 2])

(a_meanの値は [2.5, 4.0] となる)

この二つから0階成分ごとに差を取ったテンソルを求めたいのですがどのようにすればよいのか判りません
以下が求めたいテンソルです。
a[0] は a_mean[0] が引かれ a[1] は a_mean[1] の値が引かれます。

[[[-1.5, -0.5], 
  [ 0.5,  1.5]], 
 [[-3.0, -1.0], 
  [ 1.0,  3.0]]]

実際の計算では0階成分がバッチサイズとなり、上記の様な各バッチでの操作を行いたいと考えています。
ご教授お願いします。

Answer 1

a と同じ要素数を持つ a_mean を作る事で解決いたしました。

a_mean = tf.map_fn(lambda x: tf.fill(tf.shape(a)[1:3], x), a_mean)
a_dif = tf.subtract(a, a_mean)

ですが少々メモリがもったいないですね。
もっとうまいやり方がありそうな気もします。