class_docstr: Implements the differentiable sampling in mixed operation.
  
  One mixed operation can have multiple value choices in its arguments.
  Thus the ``_arch_alpha`` here is a parameter dict, and ``named_parameters``
  filters out multiple parameters with ``_arch_alpha`` as its prefix.
  
  When this class is asked for ``forward_argument``, it returns a distribution,
  i.e., a dict from int to float based on its weights.
  
  All the parameters (``_arch_alpha``, ``parameters()``, ``_softmax``) are
  saved as attributes of ``operation``, rather than ``self``,
  because this class itself is not a ``nn.Module``, and saved parameters here
  won't be optimized.
class_docstr_tok: ["Implements","the","differentiable","sampling","in","mixed","operation",".","One","mixed","operation","can","have","multiple","value","choices","in","its","arguments",".","Thus","the","``","_arch_alpha","``","here","is","a","parameter","dict",",","and","``","named_parameters","``","filters","out","multiple","parameters","with","``","_arch_alpha","``","as","its","prefix",".","When","this","class","is","asked","for","``","forward_argument","``",",","it","returns","a","distribution",",","i.e.",",","a","dict","from","int","to","float","based","on","its","weights",".","All","the","parameters","(","``","_arch_alpha","``",",","``","parameters","(",")","``",",","``","_softmax","``",")","are","saved","as","attributes","of","``","operation","``",",","rather","than","``","self","``",",","because","this","class","itself","is","not","a","``","nn.Module","``",",","and","saved","parameters","here","wo","n't","be","optimized","."]
docstr: Export is argmax for each leaf value choice.
docstr_tok: ["Export","is","argmax","for","each","leaf","value","choice","."]
code: def export(
      self, operation: MixedOperation, memo: dict[str, Any]
  ) -> dict[str, Any]:
      """Export is argmax for each leaf value choice."""
      result = {}
      for name, spec in operation.simplify().items():
          if name in memo:
              continue
          chosen_index = int(
              torch.argmax(
                  cast(dict, operation._arch_alpha)[name]
              ).item()
          )
          result[name] = cast(Categorical, spec).values[chosen_index]
      return result
  
imports: from __future__ import annotations
  import functools
  import logging
  import warnings
  from typing import Any, Dict, Sequence, List, Tuple, cast
  import torch
  import torch.nn
  import torch.nn.functional
  from nni.mutable import MutableExpression, Mutable, Categorical
  from nni.nas.nn.pytorch import LayerChoice, InputChoice, Repeat
  from nni.nas.nn.pytorch.cell import preprocess_cell_inputs
  from .base import BaseSuperNetModule
  from .operation import MixedOperation, MixedOperationSamplingPolicy
  from .sampling import PathSamplingCell
  from ._expression_utils import traverse_all_options, weighted_sum
label_desc: Use image node_id 4 for calling the MixedOpDifferentiablePolicy obj's underlying member method code with example usage: obj.export(operation, memo) and returns: result