code: def get_device_memory(
      platform_type: Optional[str], megacore: bool
  ) -> int:
      """Return the total amount of memory per device."""
      if (first_device := jax.devices()[0]).platform == "cpu":
          return 16
      elif first_device.platform == "gpu":
          if platform_type in {"p100", "v100"}:
              return 16
          elif platform_type == "p4":
              return 8
          elif platform_type == "a100":
              return 40
          else:
              return 16
      elif first_device.platform == "tpu":
          if first_device.device_kind in {
              "Cloud TPU",
              "TPU v2",
              "TPU v3",
          }:
              return 16
          elif first_device.device_kind == "TPU v4":
              if platform_type == "pl":
                  return 8
              else:
                  return 32 if megacore else 16
          else:
              raise ValueError(
                  f"Unrecognized device kind: {first_device.device_kind}"
              )
      else:
          raise ValueError(
              f"Unrecognized device type: {first_device.platform}"
          )
  
imports: from collections.abc import Callable
  from collections.abc import Iterable
  from collections.abc import Iterator
  from collections.abc import Mapping
  from collections.abc import MutableMapping
  from collections.abc import Sequence
  import copy
  import types
  from typing import Any
  from typing import Optional
  from typing import TypeVar
  from typing import Union
  from dmvr import tokenizers
  import jax
  import ml_collections
label_desc: Use image node_id 9 for calling a global function with example usage: get_device_memory(platform_type, megacore) and returns: int, int, int, int, int, int, int, unknown