class_docstr: _PyTorchReducerContext is a component of the PyTorchTrialContext which does not interact with
  the rest of the PyTorchTrialContext.  It has been separated into its own object to keep the
  size of PyTorchTrialContext down, and because there may be a future where we can create a
  _PyTorchReducerContext as a standalone feature for users of the Core API... though in that
  case we would rename it with a public name first.
class_docstr_tok: ["_PyTorchReducerContext","is","a","component","of","the","PyTorchTrialContext","which","does","not","interact","with","the","rest","of","the","PyTorchTrialContext",".","It","has","been","separated","into","its","own","object","to","keep","the","size","of","PyTorchTrialContext","down",",","and","because","there","may","be","a","future","where","we","can","create","a","_PyTorchReducerContext","as","a","standalone","feature","for","users","of","the","Core","API","...","though","in","that","case","we","would","rename","it","with","a","public","name","first","."]
docstr: Register a custom reducer that will calculate a metric properly, even with distributed
  training.
  
  During distributed training and evaluation, many types of metrics must be calculated
  globally, rather than calculating the metric on each shard of the dataset and averaged or
  summed.  For example, an accurate ROC AUC for dataset cannot be derived from the individual
  ROC AUC metrics calculated on by each worker.
  
  Determined solves this problem by offering fully customizable metric reducers which are
  distributed-aware.  These are registered by calling ``context.wrap_reducer()``
  and are updated by the user during ``train_batch()`` or ``evaluate_batch()``.
  
  Arguments:
      reducer (Union[Callable, pytorch.MetricReducer]):
          Either a reducer function or a pytorch.MetricReducer.  See below for more details.
      name: (Optional[str] = None):
          Either a string name to associate with the metric returned by the reducer, or
          ``None`` to indicate the metric will return a dict mapping string names to metric
          values.  This allows for a single reducer to return many metrics, such as for a
          per-class mean IOU calculation.  Note that if name is a string, the returned
          metric must NOT be a dict-type metric.
      for_training: (bool = True):
          Indicate that the ``reducer`` should be used for training workloads.
      for_validation: (bool = True):
          Indicate that the ``reducer`` should be used for validation workloads.
  
  Return Value:
      pytorch.MetricReducer:
          If ``reducer`` was a function, the returned ``MetricReducer`` will have a single
          user-facing method like ``def update(value: Any) -> None`` that you should call
          during ``train_batch`` or ``evaluate_batch``.  Otherwise, the return value will
          just be the ``reducer`` that was passed in.
  
  **Reducer functions: the simple API**
  
  If the ``reducer`` parameter is a function, it must have the following properties:
  
     -  It accepts a single parameter, which will be a flat list of all inputs the users
        pass when they call ``.update()`` on the object returned by ``wrap_reducer()``.
        See the code example below for more details.
     -  It returns either a single (non-dict) metric or a dictionary mapping names to
        metrics, as described above.
  
  The primary motivation for passing a function as the reducer is simplicity. Metrics from
  all batches will be buffered in memory and passed over the network before they are reduced
  all at once. This introduces some overhead, but it is likely unnoticeable for scalar
  metrics or on validation datasets of small or medium size.  This single function strategy
  can be useful for quick prototyping or for calculating metrics that are difficult
  or impossible to calculate incrementally.
  
  For example, ROC AUC could be properly calculated by passing a small wrapper function
  calling ``sklearn.metrics.roc_auc_score``:
  
  .. code-block:: python
  
     # Custom reducer function.
     def roc_auc_reducer(values):
         # values will be a flat list of all inputs to
         # .update(), which in this code example are
         # tuples of (y_true, y_score).  We reshape
         # that list into two separate lists:
         y_trues, y_scores = zip(*values)
  
         # Then we return a metric value:
         return sklearn.metrics.roc_auc_score(
             np.array(y_trues), np.array(y_scores)
         )
  
     class MyPyTorchTrial(PyTorchTrial):
         def __init__(self, context):
             self.roc_auc = context.wrap_reducer(
                 roc_auc_reducer, name="roc_auc"
             )
             ...
  
         def evaluate_batch(self, batch):
             ...
             # Function-based reducers are updated with .update().
             # The roc_auc_reducer function will get a list of all
             # inputs that we pass in here:
             self.roc_auc.update((y_true, y_score))
  
             # The "roc_auc" metric will be included in the final
             # metrics after the workload has completed; no need
             # to return it here.  If that is your only metric,
             # just return an empty dict.
             return {}
  
  **MetricReducer objects: the advanced API**
  
  The primary motivation for passing a ``det.pytorch.MetricReducer`` as the reducer is
  performance. ``det.pytorch.MetricReducer`` allows the user more control in how values are
  stored and exposes a ``per_slot_reduce()`` call which lets users minimize the cost of the
  network communication before the final ``cross_slot_reduce()``.
  
  An additional reason for using the ``det.pytorch.MetricReducer`` is for flexibility of the
  update mechanism, which is completely user-defined when subclassing ``MetricReducer``.
  
  For the full details and a code example, see: :class:`~determined.pytorch.MetricReducer`.
docstr_tok: ["Register","a","custom","reducer","that","will","calculate","a","metric","properly",",","even","with","distributed","training",".","During","distributed","training","and","evaluation",",","many","types","of","metrics","must","be","calculated","globally",",","rather","than","calculating","the","metric","on","each","shard","of","the","dataset","and","averaged","or","summed",".","For","example",",","an","accurate","ROC","AUC","for","dataset","can","not","be","derived","from","the","individual","ROC","AUC","metrics","calculated","on","by","each","worker",".","Determined","solves","this","problem","by","offering","fully","customizable","metric","reducers","which","are","distributed-aware",".","These","are","registered","by","calling","``","context.wrap_reducer","(",")","``","and","are","updated","by","the","user","during","``","train_batch","(",")","``","or","``","evaluate_batch","(",")","``",".","Arguments",":","reducer","(","Union","[","Callable",",","pytorch.MetricReducer","]",")",":","Either","a","reducer","function","or","a","pytorch.MetricReducer",".","See","below","for","more","details",".","name",":","(","Optional","[","str","]","=","None",")",":","Either","a","string","name","to","associate","with","the","metric","returned","by","the","reducer",",","or","``","None","``","to","indicate","the","metric","will","return","a","dict","mapping","string","names","to","metric","values",".","This","allows","for","a","single","reducer","to","return","many","metrics",",","such","as","for","a","per-class","mean","IOU","calculation",".","Note","that","if","name","is","a","string",",","the","returned","metric","must","NOT","be","a","dict-type","metric",".","for_training",":","(","bool","=","True",")",":","Indicate","that","the","``","reducer","``","should","be","used","for","training","workloads",".","for_validation",":","(","bool","=","True",")",":","Indicate","that","the","``","reducer","``","should","be","used","for","validation","workloads",".","Return","Value",":","pytorch.MetricReducer",":","If","``","reducer","``","was","a","function",",","the","returned","``","MetricReducer","``","will","have","a","single","user-facing","method","like","``","def","update","(","value",":","Any",")","-",">","None","``","that","you","should","call","during","``","train_batch","``","or","``","evaluate_batch","``",".","Otherwise",",","the","return","value","will","just","be","the","``","reducer","``","that","was","passed","in",".","*","*","Reducer","functions",":","the","simple","API","*","*","If","the","``","reducer","``","parameter","is","a","function",",","it","must","have","the","following","properties",":","-","It","accepts","a","single","parameter",",","which","will","be","a","flat","list","of","all","inputs","the","users","pass","when","they","call","``",".update","(",")","``","on","the","object","returned","by","``","wrap_reducer","(",")","``",".","See","the","code","example","below","for","more","details",".","-","It","returns","either","a","single","(","non-dict",")","metric","or","a","dictionary","mapping","names","to","metrics",",","as","described","above",".","The","primary","motivation","for","passing","a","function","as","the","reducer","is","simplicity",".","Metrics","from","all","batches","will","be","buffered","in","memory","and","passed","over","the","network","before","they","are","reduced","all","at","once",".","This","introduces","some","overhead",",","but","it","is","likely","unnoticeable","for","scalar","metrics","or","on","validation","datasets","of","small","or","medium","size",".","This","single","function","strategy","can","be","useful","for","quick","prototyping","or","for","calculating","metrics","that","are","difficult","or","impossible","to","calculate","incrementally",".","For","example",",","ROC","AUC","could","be","properly","calculated","by","passing","a","small","wrapper","function","calling","``","sklearn.metrics.roc_auc_score","``",":","..","code-block",":",":","python","#","Custom","reducer","function",".","def","roc_auc_reducer","(","values",")",":","#","values","will","be","a","flat","list","of","all","inputs","to","#",".update","(",")",",","which","in","this","code","example","are","#","tuples","of","(","y_true",",","y_score",")",".","We","reshape","#","that","list","into","two","separate","lists",":","y_trues",",","y_scores","=","zip","(","*","values",")","#","Then","we","return","a","metric","value",":","return","sklearn.metrics.roc_auc_score","(","np.array","(","y_trues",")",",","np.array","(","y_scores",")",")","class","MyPyTorchTrial","(","PyTorchTrial",")",":","def","__init__","(","self",",","context",")",":","self.roc_auc","=","context.wrap_reducer","(","roc_auc_reducer",",","name=","''","roc_auc","''",")","...","def","evaluate_batch","(","self",",","batch",")",":","...","#","Function-based","reducers","are","updated","with",".update","(",")",".","#","The","roc_auc_reducer","function","will","get","a","list","of","all","#","inputs","that","we","pass","in","here",":","self.roc_auc.update","(","(","y_true",",","y_score",")",")","#","The","``","roc_auc","''","metric","will","be","included","in","the","final","#","metrics","after","the","workload","has","completed",";","no","need","#","to","return","it","here",".","If","that","is","your","only","metric",",","#","just","return","an","empty","dict",".","return","{","}","*","*","MetricReducer","objects",":","the","advanced","API","*","*","The","primary","motivation","for","passing","a","``","det.pytorch.MetricReducer","``","as","the","reducer","is","performance",".","``","det.pytorch.MetricReducer","``","allows","the","user","more","control","in","how","values","are","stored","and","exposes","a","``","per_slot_reduce","(",")","``","call","which","lets","users","minimize","the","cost","of","the","network","communication","before","the","final","``","cross_slot_reduce","(",")","``",".","An","additional","reason","for","using","the","``","det.pytorch.MetricReducer","``","is","for","flexibility","of","the","update","mechanism",",","which","is","completely","user-defined","when","subclassing","``","MetricReducer","``",".","For","the","full","details","and","a","code","example",",","see",":",":","class",":","`","~determined.pytorch.MetricReducer","`","."]
code: def wrap_reducer(
      self,
      reducer: Union[Callable, MetricReducer],
      name: Optional[str] = None,
      for_training: bool = True,
      for_validation: bool = True,
  ) -> MetricReducer:
      """
      Register a custom reducer that will calculate a metric properly, even with distributed
      training.
  
      During distributed training and evaluation, many types of metrics must be calculated
      globally, rather than calculating the metric on each shard of the dataset and averaged or
      summed.  For example, an accurate ROC AUC for dataset cannot be derived from the individual
      ROC AUC metrics calculated on by each worker.
  
      Determined solves this problem by offering fully customizable metric reducers which are
      distributed-aware.  These are registered by calling ``context.wrap_reducer()``
      and are updated by the user during ``train_batch()`` or ``evaluate_batch()``.
  
      Arguments:
          reducer (Union[Callable, pytorch.MetricReducer]):
              Either a reducer function or a pytorch.MetricReducer.  See below for more details.
          name: (Optional[str] = None):
              Either a string name to associate with the metric returned by the reducer, or
              ``None`` to indicate the metric will return a dict mapping string names to metric
              values.  This allows for a single reducer to return many metrics, such as for a
              per-class mean IOU calculation.  Note that if name is a string, the returned
              metric must NOT be a dict-type metric.
          for_training: (bool = True):
              Indicate that the ``reducer`` should be used for training workloads.
          for_validation: (bool = True):
              Indicate that the ``reducer`` should be used for validation workloads.
  
      Return Value:
          pytorch.MetricReducer:
              If ``reducer`` was a function, the returned ``MetricReducer`` will have a single
              user-facing method like ``def update(value: Any) -> None`` that you should call
              during ``train_batch`` or ``evaluate_batch``.  Otherwise, the return value will
              just be the ``reducer`` that was passed in.
  
      **Reducer functions: the simple API**
  
      If the ``reducer`` parameter is a function, it must have the following properties:
  
         -  It accepts a single parameter, which will be a flat list of all inputs the users
            pass when they call ``.update()`` on the object returned by ``wrap_reducer()``.
            See the code example below for more details.
         -  It returns either a single (non-dict) metric or a dictionary mapping names to
            metrics, as described above.
  
      The primary motivation for passing a function as the reducer is simplicity. Metrics from
      all batches will be buffered in memory and passed over the network before they are reduced
      all at once. This introduces some overhead, but it is likely unnoticeable for scalar
      metrics or on validation datasets of small or medium size.  This single function strategy
      can be useful for quick prototyping or for calculating metrics that are difficult
      or impossible to calculate incrementally.
  
      For example, ROC AUC could be properly calculated by passing a small wrapper function
      calling ``sklearn.metrics.roc_auc_score``:
  
      .. code-block:: python
  
         # Custom reducer function.
         def roc_auc_reducer(values):
             # values will be a flat list of all inputs to
             # .update(), which in this code example are
             # tuples of (y_true, y_score).  We reshape
             # that list into two separate lists:
             y_trues, y_scores = zip(*values)
  
             # Then we return a metric value:
             return sklearn.metrics.roc_auc_score(
                 np.array(y_trues), np.array(y_scores)
             )
  
         class MyPyTorchTrial(PyTorchTrial):
             def __init__(self, context):
                 self.roc_auc = context.wrap_reducer(
                     roc_auc_reducer, name="roc_auc"
                 )
                 ...
  
             def evaluate_batch(self, batch):
                 ...
                 # Function-based reducers are updated with .update().
                 # The roc_auc_reducer function will get a list of all
                 # inputs that we pass in here:
                 self.roc_auc.update((y_true, y_score))
  
                 # The "roc_auc" metric will be included in the final
                 # metrics after the workload has completed; no need
                 # to return it here.  If that is your only metric,
                 # just return an empty dict.
                 return {}
  
      **MetricReducer objects: the advanced API**
  
      The primary motivation for passing a ``det.pytorch.MetricReducer`` as the reducer is
      performance. ``det.pytorch.MetricReducer`` allows the user more control in how values are
      stored and exposes a ``per_slot_reduce()`` call which lets users minimize the cost of the
      network communication before the final ``cross_slot_reduce()``.
  
      An additional reason for using the ``det.pytorch.MetricReducer`` is for flexibility of the
      update mechanism, which is completely user-defined when subclassing ``MetricReducer``.
  
      For the full details and a code example, see: :class:`~determined.pytorch.MetricReducer`.
      """
  
      # Detect double-wrapped reducers.
      if reducer in (wrapped.reducer for wrapped in self._wrapped_reducers):  # type: ignore
          raise AssertionError(
              f"Detected the same reducer of type {type(reducer).__name__} in "
              "context.wrap_reducer(), please avoid calling wrap_reducer() on the "
              "same reducer object twice."
          )
  
      if not isinstance(reducer, MetricReducer):
          if callable(reducer):
              reducer = _SimpleReducer(reducer)
          else:
              raise AssertionError(
                  f"Detected invalid reducer in wrap_reducer() of type {type(reducer).__name__}. "
                  "Reducers must either be a function or a subclass of pytorch.MetricReducer."
              )
  
      wrapped = _WrappedReducer(
          reducer, name, for_training, for_validation
      )
  
      self._wrapped_reducers.append(wrapped)
      return reducer
  
imports: import abc
  import enum
  import itertools
  from typing import Any, Callable, Dict, List, Optional, TypeVar, Union, overload
  import numpy
label_desc: Use image node_id 5 for calling the _PyTorchReducerContext obj's underlying member method code with example usage: obj.wrap_reducer(reducer, name, for_training, for_validation) and returns: reducer