code: def _estimate_memory_needs(
      inputs: list[JsonDict], edge_field: str, output_example: JsonDict
  ):
      """Estimate how much memory is needed to store all activations.
  
      We store all activations in memory to simplify the implementation, but this
      can grow to be quite large (many GB up to 100s of GB) for large encoders
      and large numbers of targets (such as OntoNotes tasks).
  
      This will log the estimated memory required, based on the dimension of
      activations and the number of total edges.
  
      Args:
        inputs: all inputs
        edge_field: name of edges field in input (such as 'coref')
        output_example: a single example from an EdgeFeaturesDataset, used to get
          embedding size
      """
      # Count total edges
      total_edges = sum([len(ex[edge_field]) for ex in inputs])
      # Get embedding size from first output example
      span1_dim = output_example["span1_embs"].size
      span2_dim = output_example["span2_embs"].size
      # pylint: disable=logging-format-interpolation
      logging.warning(
          "Found %d total edges with embedding dimensions %d + %d = %d.",
          total_edges,
          span1_dim,
          span2_dim,
          (span1_dim + span2_dim),
      )
      bytes_per_edge = (
          output_example["span1_embs"].nbytes
          + output_example["span2_embs"].nbytes
      )
      total_mbytes = total_edges * bytes_per_edge / 1e6
      if total_mbytes < 1e3:
          logging.warning(
              "Estimated memory requirement: %.3G MB", total_mbytes
          )
      else:
          logging.warning(
              "Estimated memory requirement: %.3G GB",
              (total_mbytes / 1e3),
          )
  
imports: import os
  from collections.abc import Iterable
  from typing import Optional
  from absl import logging
  from lit_nlp.api import dataset
  from lit_nlp.api import dtypes
  from lit_nlp.api import model
  from lit_nlp.api import types
  from lit_nlp.lib import utils
  import numpy
  import tensorflow
label_desc: Use image node_id 3 for calling a global function with example usage: _estimate_memory_needs(inputs, edge_field, output_example) without return types