class_docstr: Responsible for serializing GDAs via TensorStore.
class_docstr_tok: ["Responsible","for","serializing","GDAs","via","TensorStore","."]
docstr: Serializes Arrays or Arrays via TensorStore asynchronously.
  
  TensorStore writes to a storage layer in 2 steps:
  *  Reading/copying from the source after which the source can be modified.
       * Returns a copy future.
  *  Writing/committing to the storage layer.
       * Returns a commit future.
  
  In asynchronous mode, the serialization waits for the commit future to
  finish in a separate thread allowing other computation to proceed.
  
  Args:
    arrays: Arrays or Arrays that should be serialized.
    tensorstore_specs: TensorStore specs that are used to serialize GDAs or
      Arrays.
    on_commit_callback: This callback will be executed after all processes
      have finished writing their checkpoints to disk. Filesystems where
      atomic rename operations are supported, you can rename from the
      temporary directory to the final directory. On GCS, you write to the
      final directory directly and in `on_commit_callback` you write a
      success file indicating that the serialization was successful because
      GCS does not support atomic rename operations.
code: def serialize(self, arrays, tensorstore_specs, *, on_commit_callback):
      """Serializes Arrays or Arrays via TensorStore asynchronously.
  
      TensorStore writes to a storage layer in 2 steps:
      *  Reading/copying from the source after which the source can be modified.
           * Returns a copy future.
      *  Writing/committing to the storage layer.
           * Returns a commit future.
  
      In asynchronous mode, the serialization waits for the commit future to
      finish in a separate thread allowing other computation to proceed.
  
      Args:
        arrays: Arrays or Arrays that should be serialized.
        tensorstore_specs: TensorStore specs that are used to serialize GDAs or
          Arrays.
        on_commit_callback: This callback will be executed after all processes
          have finished writing their checkpoints to disk. Filesystems where
          atomic rename operations are supported, you can rename from the
          temporary directory to the final directory. On GCS, you write to the
          final directory directly and in `on_commit_callback` you write a
          success file indicating that the serialization was successful because
          GCS does not support atomic rename operations.
      """
      logger.info("Waiting for previous serialization to finish.")
      self.wait_until_finished()
  
      commit_futures = [[] for _ in range(len(tensorstore_specs))]
  
      async def _run_serializer():
          future_writer = jax.tree_util.tree_map(
              async_serialize, arrays, tensorstore_specs, commit_futures
          )
          return await asyncio.gather(*future_writer)
  
      asyncio.run(_run_serializer())
  
      self._add_futures(jax.tree_util.tree_flatten(commit_futures)[0])
  
      # Used in wait_until_finished to check on process != 0, if the checkpoint
      # has finished writing.
      self._start_async_commit(on_commit_callback)
  
imports: from __future__ import annotations
  import abc
  import asyncio
  from collections.abc import Awaitable, Sequence
  import itertools
  import logging
  from functools import partial
  import os
  import re
  import time
  import threading
  from typing import Any, Callable, Optional, Union
  import jax
  from jax._src import distributed
  from jax._src import array
  from jax._src import sharding
  from jax._src import sharding_impls
  from jax._src import typing
  from jax._src.lib import xla_extension
  import jax.numpy
  import numpy
  import tensorstore
label_desc: Use image node_id 1 for calling the GlobalAsyncCheckpointManager obj's underlying member method code with example usage: obj.serialize(arrays, tensorstore_specs) and returns: unknown