Imagine I have the following function
$$ \vec{f}(\vec{x}) = x \vec{x}, x = | \vec{x} |, \vec{x} \in R^3 $$
That is, the function is essentially a quadratic function, but contains a vector direction as well. Intuitively from single variable calculus I would expect the gradient $ \nabla \vec{f} = (\partial \vec{f}/ \partial x_1,\partial \vec{f}/ \partial x_2,\partial \vec{f}/ \partial x_3) $ to be proportional to $2x$, however I also would expect it to be a 3x3 matrix.
My most naive attempt would be to do
$$ \vec{f} = x_1^2 \vec{e}_1 + x_2^2 \vec{e}_2 + x_3^2 \vec{e}_3 $$
and say that
$$ \nabla \vec{f} = 2 x_1 \vec{e}_1 + 2 x_2 \vec{e}_2 + 2 x_3 \vec{e}_3 $$
But it would mean that every gradient w.r.t. a vector would always be a diagonal matrix, which seems wrong to me. What I really want to create is the Jacobian $ \partial \vec{f}_i / \partial x_j $ but I think I get a little bit confused about what I do with the base vectors $ \vec{e_i} $ during the partial derivative.