code: def split_for_jax_process(
      split: str,
      *,
      process_index: Optional[int] = None,
      process_count: Optional[int] = None,
      drop_remainder: bool = False,
  ) -> splits_lib.SplitArg:
      """Returns the subsplit of the data for the process.
  
      In distributed setting, all process/hosts should get a non-overlapping,
      equally sized slice of the entire data. This function takes as input a split
      and extracts the slice for the current process index.
  
      Usage:
  
      ```python
      tfds.load(..., split=tfds.split_for_jax_process('train'))
      ```
  
      This funtion is an alias for:
  
      ```python
      tfds.even_splits(split, n=jax.process_count())[jax.process_index()]
      ```
  
      By default, if examples can't be evenly distributed across processes, you can
      drop extra examples with `drop_remainder=True`.
  
      Args:
        split: Split to distribute across host (e.g. `train[75%:]`,
          `train[:800]+validation[:100]`).
        process_index: Process index in `[0, count)`. Defaults to
          `jax.process_index()`.
        process_count: Number of processes. Defaults to `jax.process_count()`.
        drop_remainder: Drop examples if the number of examples in the datasets is
          not evenly divisible by `n`. If `False`, examples are distributed evenly
          across subsplits, starting by the first. For example, if there is 11
          examples with `n=3`, splits will contain `[4, 4, 3]` examples
          respectivelly.
  
      Returns:
        subsplit: The sub-split of the given `split` for the current
          `process_index`.
      """
      if process_index is None:
          process_index = (
              lazy_imports_lib.lazy_imports.jax.process_index()
          )
      if process_count is None:
          process_count = (
              lazy_imports_lib.lazy_imports.jax.process_count()
          )
      return even_splits(
          split,
          n=process_count,
          drop_remainder=drop_remainder,
      )[process_index]
  
file_path: datasets/tensorflow_datasets/core/subsplits_utils.py
imports: import dataclasses
  import functools
  import operator
  from typing import List, Optional
  from absl import logging
  from tensorflow_datasets.core import lazy_imports_lib
  from tensorflow_datasets.core import splits
label_desc: Use image node_id 2 for calling a global function with example usage: split_for_jax_process(split) and returns: even_splits