<commit_before># Copyright 2018 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Implementation of pmap and related functionality."""

# A ShardingSpec describes at a high level how a logical array is sharded across
# devices (each ShardedDeviceArray has a ShardingSpec, and ShardingSpecs also
# describe how to shard inputs to a parallel computation). spec_to_indices()
# encodes exactly how a given ShardingSpec is translated to device buffers, i.e.
# how the sharded array is "laid out" across devices. Given a sequence of
# devices, we shard the data across the devices in row-major order, with
# replication treated as an extra inner dimension.
#
# For example, given the logical data array [1, 2, 3, 4], if we were to
# partition this array 4 ways with a replication factor of 2, for a total of 8
# devices, the data on each device would be: [1, 1], [2, 2], [3, 3], [4, 4].
#
# This encoding is assumed by various parts of the system, e.g. generating
# replica groups for collective operations.

from contextlib import contextmanager
from collections import defaultdict
import itertools as it
import operator as op
import threading
from typing import (Any, Callable, Dict, List, Optional, Sequence, Set, Tuple,
                    Type, Union, no_type_check)

from absl import logging
import numpy as np

from ..config import flags, config
from .. import core
from .. import linear_util as lu
from .. import lazy
from ..abstract_arrays import ConcreteArray, ShapedArray, array_types
from ..core import Var, Literal
from ..util import (partial, unzip2, unzip3, prod, safe_map, safe_zip,
                    extend_name_stack, wrap_name)
from ..lib import xla_bridge as xb
from ..lib import xla_client as xc
from ..tree_util import tree_flatten, tree_map
from .batching import broadcast, not_mapped, moveaxis
from . import batching
from . import partial_eval as pe
from . import xla
from . import ad


xops = xc.ops

FLAGS = flags.FLAGS

unsafe_map, map = map, safe_map

Index = Union[int, slice, Tuple[Union[int, slice], ...]]


def device_put(x, devices: Sequence[xb.xla_client.Device], replicate: bool=False) -> List[xb.xla_client._xla.PyLocalBuffer]:
  """Call device_put on a sequence of devices and return a flat sequence of buffers."""
  if replicate:
    return list(it.chain.from_iterable(xla.device_put(x, device) for device in devices))
  else:
    return list(it.chain.from_iterable(xla.device_put(val, device) for val, device in safe_zip(x, devices)))



# TODO(skye): make this a namedtuple. This may allow us to use ShardingSpecs in
# performance-sensitive code, e.g. shard_args.
class ShardingSpec:
  """Describes how a logical array is sharded across devices.

  Note this does not specify the physical devices to be sharded across, nor a
  logical ordering of data shards. Use `spec_to_indices` to resolve a
  ShardingSpec to the specific logical ordering expected throughout the system.

  Sharding includes "replication", where the same data is present on multiple
  devices. Replication is always applied to the entire logical array, i.e. the
  whole array is copied N times, although each copy may still be sharded
  according to the rest of the ShardingSpec. This means that unlike other kinds
  of sharding, replication isn't associated with a particular logical array
  axis. However, it does have a position relative to the logical array axes,
  which is necessary to specify how replication is mapped to devices in
  `spec_to_indices`. One to think about this is if you added an extra length-N
  logical axis containing the N copies of the original array, where would that
  new axis go? This would affect the final buffer order computed in
  `spec_to_indices`.

  Attributes:
    shards_per_axis: a tuple the same length as the array shape. Indicates how<commit_msg>Internal change<commit_after># Copyright 2018 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Implementation of pmap and related functionality."""

# A ShardingSpec describes at a high level how a logical array is sharded across
# devices (each ShardedDeviceArray has a ShardingSpec, and ShardingSpecs also
# describe how to shard inputs to a parallel computation). spec_to_indices()
# encodes exactly how a given ShardingSpec is translated to device buffers, i.e.
# how the sharded array is "laid out" across devices. Given a sequence of
# devices, we shard the data across the devices in row-major order, with
# replication treated as an extra inner dimension.
#
# For example, given the logical data array [1, 2, 3, 4], if we were to
# partition this array 4 ways with a replication factor of 2, for a total of 8
# devices, the data on each device would be: [1, 1], [2, 2], [3, 3], [4, 4].
#
# This encoding is assumed by various parts of the system, e.g. generating
# replica groups for collective operations.

from contextlib import contextmanager
from collections import defaultdict
import itertools as it
import operator as op
import threading
from typing import (Any, Callable, Dict, List, Optional, Sequence, Set, Tuple,
                    Type, Union, no_type_check)

from absl import logging
import numpy as np

from ..config import flags, config
from .. import core
from .. import linear_util as lu
from .. import lazy
from ..abstract_arrays import ConcreteArray, ShapedArray, array_types
from ..core import Var, Literal
from ..util import (partial, unzip2, unzip3, prod, safe_map, safe_zip,
                    extend_name_stack, wrap_name)
from ..lib import xla_bridge as xb
from ..lib import xla_client as xc
from ..tree_util import tree_flatten, tree_map
from .batching import broadcast, not_mapped, moveaxis
from . import batching
from . import partial_eval as pe
from . import xla
from . import ad


xops = xc.ops

FLAGS = flags.FLAGS

unsafe_map, map = map, safe_map

Index = Union[int, slice, Tuple[Union[int, slice], ...]]


def device_put(x, devices: Sequence[xb.xla_client.Device], replicate: bool=False) -> List[xb.xla_client._xla.PyLocalBuffer]:
  """Call device_put on a sequence of devices and return a flat sequence of buffers."""
  if replicate:
    return list(it.chain.from_iterable(xla.device_put(x, device) for device in devices))
  else:
    return list(it.chain.from_iterable(xla.device_put(val, device) for val, device in safe_zip(x, devices)))



# TODO(skye): make this a namedtuple. This may allow us to use ShardingSpecs in
# performance-sensitive code, e.g. shard_args.
class ShardingSpec:
  """Describes how a logical array is sharded across devices.

  Note this does not specify the physical devices to be sharded across, nor a
  logical ordering of data shards. Use `spec_to_indices` to resolve a
  ShardingSpec to the specific logical ordering expected throughout the system.

  Sharding includes "replication", where the same data is present on multiple
  devices. Replication is always applied to the entire logical array, i.e. the
  whole array is copied N times, although each copy may still be sharded
  according to the rest of the ShardingSpec. This means that unlike other kinds
  of sharding, replication isn't associated with a particular logical array
  axis. However, it does have a position relative to the logical array axes,
  which is necessary to specify how replication is mapped to devices in
  `spec_to_indices`. One way to think about this is if you added an extra
  length-N logical axis containing the N copies of the original array, where
  would that new axis go? This would affect the final buffer order computed in
  `spec_to_indices`.

  Attributes:
    shards_per_axis: a tuple the same length as the array shape. Indicates how
{'fraction_non_alphanumeric': 0.06492000927428704, 'fraction_numerical': 0.005796429399489914, 'mean_word_length': 4.187612748045701, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3185770', 'n_tokens_mistral': 2391, 'n_tokens_neox': 2313, 'n_words': 1193}