code: def compute_flops(
      flax_model_apply_fn: Callable[[jnp.ndarray], Any],
      input_spec: Mapping[str, Tuple[Tuple[int, ...], jnp.dtype]],
      fuse_multiply_add: bool,
  ) -> float:
      """Performs static analysis of the graph to compute theoretical FLOPs.
  
      This function is branched from scenic/common_lib/debug_utils.py.
      The difference is that here the input_spec is a dictionary while it is a
      sequence in the original implementation.
  
      Args:
        flax_model_apply_fn: Apply function of the flax model to be analysed.
        input_spec: An mapping of modality names to (shape, dtype) pairs specifying
          the shape and dtype of the inputs.
        fuse_multiply_add: Bool; If true, count a multiply and add (also known as
          "multiply-accumulate" or "MAC") as 1 FLOP rather than 2 (as done by the
          HLO analysis). This is commonly used in literature.
  
      Returns:
        flops: The total number of flops.
      """
      input_placeholder = {}
      for modality, spec in input_spec.items():
          in_st = debug_utils.input_spec_to_jax_shape_dtype_struct(
              spec, batch_size=1
          )
          input_placeholder[modality] = jnp.zeros(
              in_st.shape, in_st.dtype
          )
  
      analysis = (
          jax.jit(flax_model_apply_fn)
          .lower(input_placeholder)
          .cost_analysis()
      )
  
      flops = analysis["flops"]
      if fuse_multiply_add:
          flops = flops / 2
      logging.info(
          "GFLOPs %0.3f for input spec: %s", flops / 10**9, input_spec
      )
      return flops
  
file_path: scenic/scenic/projects/mbt/train_utils.py
imports: import functools
  from typing import Any, Callable, Tuple, Optional, Mapping, Union
  from absl import logging
  import flax
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.common_lib import debug_utils
  from scenic.train_lib_deprecated import optimizers
label_desc: Use image node_id 1 for calling a global function with example usage: compute_flops(flax_model_apply_fn, input_spec, fuse_multiply_add) and returns: flops