def make_output_postprocess_fn(
      cls,
      config: ml_collections.ConfigDict  # pylint: disable=unused-argument
  ) -> Callable[[Dict[str, Any], Dict[str, Any]], Dict[str, Any]]:
    """Produces function to postprocess task samples (input and output).

    The method is occasionally called during training or evaluation to save
    model inputs and outputs for manual inspection. Given an input batch and
    model's auxiliary output, the method needs to produce JSON-serializable
    dictionary with all the relevant features.

    This default implementation provides a function that takes all of the
    features in the `batch` and in `auxiliary_output`, adjust array types if
    necessary (jnp.bfloat16 is not JSON-serializable) and converts all JAX and
    numpy arrays to list (arrays are not JSON-serializable).

    Args:
      config: contains experiment hyperparameters.

    Returns:
      Function that postprocesses model's input and output for serialization.
    """

    def postprocess_fn(batch: Dict[str, Any],
                       auxiliary_output: Dict[str, Any]) -> Dict[str, Any]:
      """Function that prepares model's input and output for serialization."""

      features = {}
      for dict_of_features in [batch, auxiliary_output]:
        for key, value in dict_of_features.items():
          convert_to_list = (
              isinstance(value, Array) or isinstance(value, np.ndarray))
          if convert_to_list and value.dtype == jnp.bfloat16:
            # bfloat16 is not JSON serializable => convert to float32.
            features[key] = value.astype(jnp.float32).tolist()
          elif convert_to_list:
            # Cannot serialize numpy / jax arrays to json => convert to list.
            features[key] = value.tolist()
          else:
            features[key] = value

      return features

    return postprocess_fn
Produces function to postprocess task samples (input and output).

    The method is occasionally called during training or evaluation to save
    model inputs and outputs for manual inspection. Given an input batch and
    model's auxiliary output, the method needs to produce JSON-serializable
    dictionary with all the relevant features.

    This default implementation provides a function that takes all of the
    features in the `batch` and in `auxiliary_output`, adjust array types if
    necessary (jnp.bfloat16 is not JSON-serializable) and converts all JAX and
    numpy arrays to list (arrays are not JSON-serializable).

    Args:
      config: contains experiment hyperparameters.

    Returns:
      Function that postprocesses model's input and output for serialization.
    
language/mentionmemory/tasks/base_task.py