class_docstr: Compute the gradient of ``arch_alpha``,
  given the sampled index, inputs and outputs of the layer, and a function to
  compute the output of the layer conditioned on any other index.
  
  The forward of this function is merely returning a "copy" of the given input.
  This is different from the
  `official implementation <https://github.com/mit-han-lab/proxylessnas/blob/9cdd0791/search/modules/mix_op.py>`,
  where the output is computed within this function.
  
  Things tried but NOT WORKING:
  
  1. Use ``full_backward_hook`` instead of ``autograd.Function``.
     Since ``full_backward_hook`` is not intended for computing gradients for parameters,
     The gradients get overridden when another loss is added (e.g., latency loss).
  2. Computing the output in ``autograd.Function`` like the official impl does.
     This requires calling either ``autograd.grad`` (doesn't work when nested),
     or ``autograd.backward`` in the backward of the Function class.
     The gradients within the candidates (like Linear, or Conv2d) will computed inside this function.
     This is problematic with DDP, because either (i) with ``find_unused_parameters=True``,
     the parameters within the candidates are considered unused (because they are not found on the autograd graph),
     and raises error when they receive a gradient (something like "another gradient received after ready"), or
     (ii) with ``find_unused_parameters=False``, DDP complains about some candidate paths having no gradients at all.
  
  Therefore, the only way to make it work is to write a function::
  
      func(alpha, input, output) = output
  
  From the outside of this box, ``func`` echoes the output. But inside, there is gradient going on.
  When back-propagation, ``alpha`` will receive a computed gradient.
  ``input`` receives none because it didn't participate in the computation
  (but only needs to be saved for computing gradients).
  ``output`` receives the gradient that ``func`` receives.
class_docstr_tok: ["Compute","the","gradient","of","``","arch_alpha","``",",","given","the","sampled","index",",","inputs","and","outputs","of","the","layer",",","and","a","function","to","compute","the","output","of","the","layer","conditioned","on","any","other","index",".","The","forward","of","this","function","is","merely","returning","a","``","copy","''","of","the","given","input",".","This","is","different","from","the","`","official","implementation","<","https",":","\/\/github.com\/mit-han-lab\/proxylessnas\/blob\/9cdd0791\/search\/modules\/mix_op.py",">","`",",","where","the","output","is","computed","within","this","function",".","Things","tried","but","NOT","WORKING",":","1",".","Use","``","full_backward_hook","``","instead","of","``","autograd.Function","``",".","Since","``","full_backward_hook","``","is","not","intended","for","computing","gradients","for","parameters",",","The","gradients","get","overridden","when","another","loss","is","added","(","e.g.",",","latency","loss",")",".","2",".","Computing","the","output","in","``","autograd.Function","``","like","the","official","impl","does",".","This","requires","calling","either","``","autograd.grad","``","(","does","n't","work","when","nested",")",",","or","``","autograd.backward","``","in","the","backward","of","the","Function","class",".","The","gradients","within","the","candidates","(","like","Linear",",","or","Conv2d",")","will","computed","inside","this","function",".","This","is","problematic","with","DDP",",","because","either","(","i",")","with","``","find_unused_parameters=True","``",",","the","parameters","within","the","candidates","are","considered","unused","(","because","they","are","not","found","on","the","autograd","graph",")",",","and","raises","error","when","they","receive","a","gradient","(","something","like","``","another","gradient","received","after","ready","''",")",",","or","(","ii",")","with","``","find_unused_parameters=False","``",",","DDP","complains","about","some","candidate","paths","having","no","gradients","at","all",".","Therefore",",","the","only","way","to","make","it","work","is","to","write","a","function",":",":","func","(","alpha",",","input",",","output",")","=","output","From","the","outside","of","this","box",",","``","func","``","echoes","the","output",".","But","inside",",","there","is","gradient","going","on",".","When","back-propagation",",","``","alpha","``","will","receive","a","computed","gradient",".","``","input","``","receives","none","because","it","did","n't","participate","in","the","computation","(","but","only","needs","to","be","saved","for","computing","gradients",")",".","``","output","``","receives","the","gradient","that","``","func","``","receives","."]
code: def forward(
      ctx,
      forward_path_func,
      sample_index,
      num_inputs,
      softmax,
      arch_alpha,
      *layer_input_output,
  ):
      ctx.forward_path_func = forward_path_func
      ctx.sample_index = sample_index
      # First num_inputs are inputs, the rest of outputs.
      ctx.num_inputs = num_inputs
      ctx.softmax = softmax
  
      ctx.save_for_backward(arch_alpha, *layer_input_output)
  
      layer_output = layer_input_output[num_inputs:]
  
      # Why requires_grad: This function can be considered as an operator of arch_alpha + some inputs.
      #                    It is a differentiable function. The results must require grad.
      # Why detach: So that requires_grad flag is not directly put onto input.
      # Why unpack_tuple: Output is packed to tuple when sending to this forward.
      #                   Need to restore the original object structure.
      # NOTE: This could be potentially troublesome if users actually returned a single-ele tuple.
      return _unpack_tuple(
          _detach_tensor(layer_output, requires_grad=True)
      )
  
imports: from __future__ import annotations
  from typing import Any, cast
  import torch
  import torch.nn
  from nni.mutable import Categorical, Mutable, Sample, SampleValidationError, MutableExpression, label_scope, auto_label
  from nni.nas.nn.pytorch import Repeat, recursive_freeze
  from .base import BaseSuperNetModule
  from .differentiable import DifferentiableMixedLayer, DifferentiableMixedInput
  from ._expression_utils import traverse_all_options
label_desc: Use image node_id 1 for calling the _ProxylessFunction obj's underlying member method code with example usage: obj.forward(ctx, forward_path_func, sample_index, num_inputs, softmax, arch_alpha) and returns: _unpack_tuple