code: def compute_flops(
      flax_model_apply_fn: Callable[[jnp.ndarray], Any],
      input_spec: Sequence[
          Union[
              Tuple[Tuple[int, ...], jnp.dtype], Tuple[int, ...], None
          ]
      ],
      fuse_multiply_add: bool,
  ) -> float:
      """Performs static analysis of the graph to compute theoretical FLOPs.
  
      One can also use the XProf profiler to get the actual FLOPs at runtime
      based on device counters. Theoretical FLOPs are more useful for comparing
      models across different library implementations and is hardware-agnostic.
  
      Args:
        flax_model_apply_fn: Apply function of the flax model to be analysed.
        input_spec: An iterable of (shape, dtype) pairs specifying the shape and
          dtype of the inputs. If unspecified the dtype is float32.
        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.
      """
      dummy_input = []
      for spec in input_spec:
          if spec is not None:
              in_st = input_spec_to_jax_shape_dtype_struct(
                  spec, batch_size=1
              )
              dummy_input.append(jnp.zeros(in_st.shape, in_st.dtype))
          else:
              dummy_input.append(None)
  
      analysis = (
          jax.jit(flax_model_apply_fn)
          .lower(*dummy_input)
          .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
  
imports: from collections import abc
  from concurrent import futures
  import json
  import operator
  import threading
  from typing import Any, Callable, Optional, Sequence, Set, Tuple, Union
  from absl import logging
  from clu import parameter_overview
  import jax
  import jax.numpy
  from jax.tree_util import tree_map
  import ml_collections
label_desc: Use image node_id 4 for calling a global function with example usage: compute_flops(flax_model_apply_fn, input_spec, fuse_multiply_add) and returns: flops