code: def _get_flops(fn, *args, **kwargs):
      e = jax.jit(fn).lower(*args, **kwargs)
      cost = e.cost_analysis()
      flops = int(cost["flops"]) if "flops" in cost else 0
      return flops
  
imports: import dataclasses
  import enum
  import io
  from abc import ABC, abstractmethod
  from types import MappingProxyType
  from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Set, Tuple, Type, Union
  import jax
  import jax.numpy
  import numpy
  import rich.console
  import rich.table
  import rich.text
  import yaml
  import flax.linen.module
  from flax.core import meta, unfreeze
  from flax.core.scope import CollectionFilter, DenyList, FrozenVariableDict, LazyRng, MutableVariableDict
label_desc: Use image node_id 2 for calling a global function with example usage: _get_flops(fn) and returns: flops