tensorflow - TensorFlow Eager Execution APIでベクトルごとにヤコビ行列形式の導関数を取得するには?

Question

MLP モデルでは、層 l の入力は次の式で計算できます: z = Wa + b W は層l-1と層の間の重み行列、a は層ニューロンlの出力信号、b は層のバイアスです。例えば：l-1l

TensorFlow Eager Execution API を使用して派生物を取得したいと考えています。

z の値を計算する関数を定義します。

def f002(W, a, b):
    return tf.matmul(W, a) + b

私の主なプログラム:

def test001(args={}):
    tf.enable_eager_execution()
    tfe = tf.contrib.eager

    a = tf.reshape(tf.constant([1.0, 2.0, 3.0]), [3, 1])
    W = tf.constant([[4.0, 5.0, 6.0],[7.0, 8.0, 9.0]])
    b = tf.reshape(tf.constant([1001.0, 1002.0]), [2, 1])
    z = f002(W, a, b)
    print(z)
    grad_f1 = tfe.gradients_function(f002)
    dv = grad_f1(W, a, b)
    print(dv)

フォワード モードで z の正しい値を取得できます。しかし、派生結果を印刷すると、次のように表示されました。

[<tf.Tensor: id=17, shape=(2, 3), dtype=float32, numpy=
array([[1., 2., 3.],
       [1., 2., 3.]], dtype=float32)>, <tf.Tensor: id=18, shape=(3, 1), 
dtype=float32, numpy=
array([[11.],
       [13.],
       [15.]], dtype=float32)>, <tf.Tensor: id=16, shape=(2, 1), 
dtype=float32, numpy=
array([[1.],
       [1.]], dtype=float32)>]

これは私が望むものではありません。ベクトルごとのヤコビ行列微分結果を取得する方法は?