code: def traverse_all_options(
      mutable_expr: MutableExpression[T],
      weights: dict[str, Sequence[float]]
      | dict[str, list[float]]
      | dict[str, np.ndarray]
      | dict[str, torch.Tensor]
      | None = None,
  ) -> list[tuple[T, float]] | list[T]:
      """Traverse all possible computation outcome of a value choice.
      If ``weights`` is not None, it will also compute the probability of each possible outcome.
  
      NOTE: This function is very similar to ``MutableExpression.grid``,
      but it supports specifying weights for each leaf node.
  
      Parameters
      ----------
      mutable_expr
          The value choice to traverse.
      weights
          If there's a prior on leaf nodes, and we intend to know the (joint) prior on results,
          weights can be provided. The key is label, value are list of float indicating probability.
          Normally, they should sum up to 1, but we will not check them in this function.
  
      Returns
      -------
      Results will be sorted and duplicates will be eliminated.
      If weights is provided, the return value will be a list of tuple, with option and its weight.
      Otherwise, it will be a list of options.
      """
  
      # Validation
      simplified = mutable_expr.simplify()
      for label, param in simplified.items():
          if not isinstance(param, Categorical):
              raise TypeError(
                  f"{param!r} is not a categorical distribution"
              )
          if weights is not None:
              if label not in weights:
                  raise KeyError(
                      f"{mutable_expr} depends on a weight with key {label}, but not found in {weights}"
                  )
              if len(param) != len(weights[label]):
                  raise KeyError(
                      f"Expect weights with {label} to be of length {len(param)}, but {len(weights[label])} found"
                  )
  
      # result is a dict from a option to its weight
      result: dict[T, float] = {}
  
      sample = {}
      for sample_res in mutable_expr.grid(memo=sample):
          probability = 1.0
          if weights is not None:
              for label, chosen in sample.items():
                  if isinstance(weights[label], dict):
                      # weights[label] is a dict. Choices are used as keys.
                      probability = probability * weights[label][chosen]
                  else:
                      # weights[label] is a list. We need to find the index of currently chosen value.
                      chosen_idx = cast(
                          Categorical, simplified[label]
                      ).values.index(chosen)
                      if chosen_idx == -1:
                          raise RuntimeError(
                              f"{chosen} is not a valid value for {label}: {simplified[label]!r}"
                          )
                      probability = (
                          probability * weights[label][chosen_idx]
                      )
  
          if sample_res in result:
              result[sample_res] = result[sample_res] + cast(
                  float, probability
              )
          else:
              result[sample_res] = cast(float, probability)
  
      if weights is None:
          return sorted(result.keys())  # type: ignore
      else:
          return sorted(result.items())
  
file_path: nni/nni/nas/oneshot/pytorch/supermodule/_expression_utils.py
imports: from __future__ import annotations
  import operator
  from typing import Any, TypeVar, List, cast, Mapping, Sequence, Optional, Iterable, overload
  import numpy
  import torch
  from nni.mutable import MutableExpression, Categorical
label_desc: Use image node_id 4 for calling a global function with example usage: traverse_all_options(mutable_expr, weights) and returns: sorted, sorted