dlgradient of a subset of variables

I would like to apply dlgradient to a parametrized function, which would allow me in some cases to get the complete gradient of the function, but in others would take only the derivative with respect to the i-th variable. My attempt below fails. Is there a way that both succeeds and which would be efficient in the more general case when numel(x0) is large?

x0 = dlarray( [1; 2] );

i=[];

y_grad=dlfeval(@(x)oneDeriv(x,i), x0)  %works

y_grad = 
  2×1 dlarray

     1
     2

i=2;

dy_dxi=dlfeval(@(x)oneDeriv(x,i), x0) %fails

dy_dxi = 
  1×1 dlarray

     0

function  out = oneDeriv(x,i)
   
    y=sum(x.^2)/2;

    if isempty(i) %take complete gradient

       out=dlgradient(y,x); 

    else %take gradient only w.r.t x(i)

       out=dlgradient(y,x(i)); 

    end
  
   
end

Answer 1

[Matt J]

This seems to work:

X0 = dlarray( [1; 2; 3; 4]*10 );

subset=[2,3];

gradTotal = getGradient(X0) %total gradient

gradTotal = 
  4×1 dlarray

    10
    20
    30
    40

gradSubset = getGradient(X0,2:3) %gradient on subset of x

gradSubset = 
  2×1 dlarray

    20
    30

function grad = getGradient(Xall,subset)

  if nargin<2, subset=':'; end

  grad  = dlfeval( @(xsub)theFunction(X,xsub,subset), Xall(subset));

end

function  grad = theFunction(Xall,xsub,subset)
   
   Xall(subset)=xsub;

   y=sum(Xall.^2)/2;

   grad=dlgradient(y,xsub);

end

Answer 2

[Ben]

This is a subtle part of the dlarray autodiff system, the line dlgradient(y,x(i)) returns 0 because it sees the operation x -> x(i), and only knows that y depends on x, not x(i).

You can work around this by computing all the derivatives and do indexing afterwards:

function out = oneDeriv(x,i)

y = sum(x.^2)/2;

out = dlgradient(y,x);

if ~isempty(i)

out = out(i);

end

end

[Matt J]

Thanks, Ben. So the answer is no?