'Build the graph for the per-example gradient through the op.
Assumes that the MatMul was called with a design matrix with examples
in rows as the first argument and parameters as the second argument.
Args:
x: The Tensor to differentiate with respect to. This tensor must
represent the weights.
z_grads: The list of gradients on the output of the op.
Returns:
x_grads: A Tensor containing the gradient with respect to `x` for
each example. This is a 3-D tensor, with the first axis corresponding
to examples and the remaining axes matching the shape of x.'

def __call__(self, x, z_grads):

    idx = list(self.op.inputs).index(x)
    assert (idx != (-1))
    assert (len(z_grads) == len(self.op.outputs))
    assert (idx == 1)
    (x, _) = self.op.inputs
    (z_grads,) = z_grads
    x_expanded = tf.expand_dims(x, 2)
    z_grads_expanded = tf.expand_dims(z_grads, 1)
    return tf.multiply(x_expanded, z_grads_expanded)
