Obtaining the diagonal of Jacobians inside a piecewise function

Question

I have a piecewise function that has in each case an N number of functions that take N arguments. I want to retrieve the diagonals of the Jacobian, or simply the element-wise first derivative of the N functions and N arguments.

Example code:

i[a_,b_]:=Piecewise[
  {{{f[a,b],g[a,b]},a>b},
  {{j[a,b],k[a,b]},2a<b}},
  {l[a,b],m[a,b]}
]
D[i[a,b],{{a,b}}]

D can handle piecewise functions and provides the full Jacobians inside of the piecewise function.

How do I obtain a piecewise function that contains the diagonals of the Jacobian instead of the full Jacobian?

Can I rewrite this so Mathematica does not calculate the full Jacobian, which when N is large results in wasteful unnecessary computation?

EDIT In looking at MichealE2's helpful answer to this question, I was trying to think of a more general way to apply functions to expressions inside piecewise functions and found this kglr piecewise function mapping solution. Reproduced and and applied with MichaelE2's MapThread approach here:

pwMap[f_] := MapAt[f,#,{{1, All, 1},{2}}]&;
i[a,b]//pwMap[MapThread[D,{#,{a,b}}]&]

This is nice because it should be a completely general way to apply functions to expressions inside piecewise functions.

Answer 1

Instead of creating the whole Jacobian and then taking just the diagonal, we can just create the elements on their own with Inner. And we pick out the functions in Piecewise with the pattern funcs_List /; Depth[funcs] == 3 and replace them with the diagonal of their associated Jacobian.

i[x, y] /. funcs_List /; Depth[funcs] == 3 :> Inner[D, funcs, {x, y}, List]

$\begin{cases} \left\{f^{(1,0)}(x,y),g^{(0,1)}(x,y)\right\} & x>y \\ \left\{j^{(1,0)}(x,y),k^{(0,1)}(x,y)\right\} & 2 x<y \\ \left\{ \left\{ \begin{array}{cc} l^{(1,0)}(x,y) & m^{(0,1)}(x,y) \\ \end{array} \right\} \right\} & \text{True} \end{cases}$

Answer 2

First Q:

PiecewiseExpand@ Quiet@ Diagonal[D[i[a, b], {{a, b}}]]

Second Q:

diagJac[f_, vars_?VectorQ] := Module[{dmap},
   dmap[ff_?VectorQ] := MapThread[D, {ff, vars}];
   PiecewiseExpand@ dmap@ f];

diagJac[i[a, b], {a, b}]