| # Copyright 2017 The TensorFlow Authors. All Rights Reserved. |
| # |
| # 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 |
| # |
| # http://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. |
| # ============================================================================= |
| """Helper library for sharding during TPU compilation.""" |
| |
| |
| from tensorflow.python.framework import tensor_shape |
| |
| _DEFAULT_NUMBER_OF_SHARDS = 1 |
| _DEFAULT_SHARD_DIMENSION = 0 |
| |
| |
| # TODO(b/36777903) change other parts of tpu.py to use this class. |
| class ShardingPolicy(object): |
| """An object use to hold the sharding policy for a Tensor.""" |
| |
| def __init__(self): |
| self._number_of_shards = None |
| self._number_of_partitions = 1 |
| self._shard_dimension = None |
| self._frozen = False |
| |
| def __str__(self): |
| if self.number_of_shards is None or self.shard_dimension is None: |
| return "ShardingPolicy(unset)" |
| else: |
| return ("ShardingPolicy(%d shards dimension %d)" % |
| (self.number_of_shards, self.shard_dimension)) |
| |
| def _fill_default_values(self): |
| if self._number_of_shards is None: |
| self._number_of_shards = _DEFAULT_NUMBER_OF_SHARDS |
| if self._shard_dimension is None: |
| self._shard_dimension = tensor_shape.as_dimension( |
| _DEFAULT_SHARD_DIMENSION) |
| |
| def freeze(self): |
| """Prevents further modification to the sharding policy. |
| |
| Any values that have not been set when freeze is called are set to |
| defaults. If the ShardingPolicy is already frozen, this is a NoOp. |
| """ |
| if not self._frozen: |
| self._fill_default_values() |
| self._frozen = True |
| |
| @property |
| def number_of_shards(self): |
| """Returns the number of shards in the policy or None if unspecified.""" |
| return self._number_of_shards |
| |
| def set_number_of_shards(self, number_of_shards): |
| """Sets the number of shards for the current policy. |
| |
| If the policy has been frozen then number_of_shards must match the |
| existing setting. |
| |
| Args: |
| number_of_shards: The number of shards to use in the policy. |
| |
| Raises: |
| ValueError: If the policy has been frozen and number_of_shards |
| differs from the frozen value; or number_of_shards <= 0. |
| """ |
| if self._frozen: |
| if self._number_of_shards != number_of_shards: |
| raise ValueError( |
| f"Can't set sharding policy to use {number_of_shards} shards since " |
| f"it has been frozen to use {self._number_of_shards}") |
| else: |
| if number_of_shards > 0: |
| self._number_of_shards = number_of_shards |
| else: |
| raise ValueError( |
| f"Can't set sharding policy to use {number_of_shards} shards; " |
| "value must be > 0") |
| |
| @property |
| def number_of_partitions(self): |
| """Returns the number of partitions of the policy or None if unspecified.""" |
| return self._number_of_partitions |
| |
| def set_number_of_partitions(self, number_of_partitions): |
| """Sets the number of partitions for the current policy. |
| |
| If the policy has been frozen then shard_dimension must match the |
| existing setting. |
| |
| Args: |
| number_of_partitions: The number of partitions to use in the policy. |
| |
| Raises: |
| ValueError: If the policy has been frozen and shard_dimension |
| differs from the frozen value. |
| """ |
| if self._frozen: |
| if self._number_of_partitions != number_of_partitions: |
| raise ValueError( |
| f"Can't set number_of_partitions to {number_of_partitions} since " |
| f"it has been frozen to use {self._number_of_partitions}.") |
| else: |
| self._number_of_partitions = number_of_partitions |
| |
| @property |
| def shard_dimension(self): |
| """Returns the shard dimension of the policy or None if unspecified.""" |
| return self._shard_dimension |
| |
| def set_shard_dimension(self, shard_dimension): |
| """Sets the shard dimension for the current policy. |
| |
| If the policy has been frozen then shard_dimension must match the |
| existing setting. |
| |
| Args: |
| shard_dimension: The shard dimension to use in the policy. |
| |
| Raises: |
| ValueError: If the policy has been frozen and shard_dimension |
| differs from the frozen value, or shard_dimension can't be |
| interpreted as a Dimension. |
| """ |
| if self._frozen: |
| if self._shard_dimension != shard_dimension: |
| raise ValueError( |
| "Can't set shard dimension to %d since it has been frozen to " |
| "use %d." % (shard_dimension, self._shard_dimension)) |
| else: |
| self._shard_dimension = tensor_shape.as_dimension(shard_dimension) |
| |
| def merge(self, other): |
| """Merges the policy of another policy into the current policy. |
| |
| Args: |
| other: The policy to merge into this one. |
| |
| Raises: |
| ValueError: If this policy has been frozen and the merge conflicts with |
| the frozen policy. |
| """ |
| if other.number_of_shards is not None: |
| self.set_number_of_shards(other.number_of_shards) |
| if other.shard_dimension is not None: |
| self.set_shard_dimension(other.shard_dimension) |
| |
| def get_unpartitioned_shape(self, shape): |
| """Returns the shape of an unpartitioned Tensor. |
| |
| When given the shape of a 'sharded-size' Tensor, returns the shape |
| of the full shape of its unpartitioned Tensor. |
| |
| Args: |
| shape: The shape of the sharded Tensor. |
| |
| Returns: |
| The shape of the unpartitioned version of the Tensor. |
| |
| Raises: |
| ValueError: if shape has unknown sharded dimension |
| """ |
| shape = tensor_shape.as_shape(shape) |
| dims = shape.as_list() |
| if (self._shard_dimension is None or self._number_of_partitions is None or |
| not dims): |
| return None |
| if dims[self._shard_dimension] is None: |
| raise ValueError(f"Shape {shape.as_list()} must have a fixed size for " |
| f"dimension {self._shard_dimension} that is known. ") |
| if self._number_of_partitions > 1: |
| dims[self._shard_dimension] *= self._number_of_partitions |
| return tensor_shape.as_shape(dims) |
| |
| def get_sharded_shape(self, shape, shard_index=None): |
| """Returns the shape of a shard of a full Tensor. |
| |
| When given the shape of a 'full-size' Tensor, returns the shape of |
| the sub-Tensor after it has been sharded. Freezes the policy if it |
| has not yet been frozen. |
| |
| Args: |
| shape: The shape of the full-size Tensor to be sharded. |
| shard_index: The index of the shard whose shape should be returned. |
| shard_index can be None for sharding policies that use the same shape |
| for every shard. |
| |
| Returns: |
| The shape of the sharded version of the Tensor. |
| |
| Raises: |
| ValueError: If shard_index is None when shards are of different |
| shapes; or shard_index is not None and |
| !(0<=shard_index<number_of_shards); or shape does not have at |
| least self.shard_dimension+1 dimensions; or the value of |
| shape's shard dimension is not a multiple of |
| self.number_of_shards |
| """ |
| if self._shard_dimension is None or self._number_of_shards is None: |
| # Don't raise an error if the config is unset. |
| return None |
| if shard_index is not None: |
| if shard_index < 0 or shard_index >= self.number_of_shards: |
| raise ValueError( |
| f"Requested shard_index {shard_index}, but shard_index must be in " |
| f"[0,{self._number_of_shards}).") |
| shape = tensor_shape.as_shape(shape) |
| if self._number_of_shards == 1: |
| # Don't do anything when there's only one shard. |
| return shape |
| ndims = shape.ndims |
| if ndims is None: |
| raise ValueError(f"Shape {shape} must be a known shape.") |
| if ndims <= self._shard_dimension: |
| raise ValueError( |
| f"Shape {shape.as_list()} does not contain shard_dimension " |
| f"{self._shard_dimension}") |
| dims = shape.as_list() |
| if dims[self._shard_dimension] is None: |
| raise ValueError( |
| f"Shape {shape.as_list()} must have a fixed size for dimension " |
| f"{self._shard_dimension} that is known at construction time.") |
| if (dims[self._shard_dimension] % self._number_of_shards) != 0: |
| raise ValueError( |
| f"Shape {shape.as_list()} cannot be sharded {self._number_of_shards} " |
| f"ways along dimension {self._shard_dimension}") |
| dims[self._shard_dimension] //= self._number_of_shards |
| return tensor_shape.TensorShape(dims) |
| |
| def _unshard_shape(self, shape): |
| """Return the unsharded shape that would generate a given sharded shape. |
| |
| Args: |
| shape: the sharded shape to unshard |
| |
| Returns: |
| The unsharded shape. |
| |
| Raises: |
| ValueError: if shape is unknown or does not contain |
| self.shard_dimension |
| TypeError: if shape is not convertible to a TensorShape |
| """ |
| shape = tensor_shape.as_shape(shape) |
| if self._number_of_shards == 1: |
| # Don't do anything when there's only one shard. |
| return shape |
| ndims = shape.ndims |
| if ndims is None: |
| raise ValueError(f"Shape {shape} must be statically known.") |
| if ndims <= self._shard_dimension: |
| raise ValueError(f"Shape {shape.as_list()} does not contain " |
| f"shard_dimension {self._shard_dimension}. " |
| f"Rank is too small.") |
| dims = shape.as_list() |
| dims[self._shard_dimension] *= self._number_of_shards |
| return tensor_shape.TensorShape(dims) |
| |
| def get_unsharded_shape(self, shapes): |
| """Returns the shape of an unsharded Tensor given a list of shards. |
| |
| When given a list of shapes of shards, returns the shape of the |
| unsharded Tensor that would generate the shards. Sets defaults for the |
| policy if number_of_shards or shard_dimension is None. |
| |
| Args: |
| shapes: The shapes of the Tensor shards to be combined. |
| |
| Returns: |
| The shape of the unsharded version of the Tensor. |
| |
| Raises: |
| ValueError: if shapes is not a list of length |
| self.number_of_shards; or any element of shapes is not a valid |
| shape consistent with the sharding policy; or the list of |
| shapes is not a valid sharding of a full shape. |
| TypeError: if an element of shapes is not convertible to a |
| TensorShape |
| """ |
| self._fill_default_values() |
| if len(shapes) != self.number_of_shards: |
| raise ValueError( |
| f"Shapes {shapes} is length {len(shapes)} but must be a list of " |
| f"length number_of_shards={self.number_of_shards}") |
| unsharded_shapes = [self._unshard_shape(s) for s in shapes] |
| for i in range(self.number_of_shards - 1): |
| if not unsharded_shapes[i].is_compatible_with( |
| unsharded_shapes[self.number_of_shards - 1]): |
| raise ValueError( |
| f"Sharded shapes {shapes} are not consistent shards of a full shape " |
| f"sharded {self.number_of_shards} ways along " |
| f"dimension {self.shard_dimension}.") |
| return unsharded_shapes[0] |