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. sh
(self, shape, shard_index=None)
| 135 | self.set_shard_dimension(other.shard_dimension) |
| 136 | |
| 137 | def get_sharded_shape(self, shape, shard_index=None): |
| 138 | """Returns the shape of a shard of a full Tensor. |
| 139 | |
| 140 | When given the shape of a 'full-size' Tensor, returns the shape of |
| 141 | the sub-Tensor after it has been sharded. Freezes the policy if it |
| 142 | has not yet been frozen. |
| 143 | |
| 144 | Args: |
| 145 | shape: The shape of the full-size Tensor to be sharded. |
| 146 | shard_index: The index of the shard whose shape should be returned. |
| 147 | shard_index can be None for sharding policies that use the same |
| 148 | shape for every shard. |
| 149 | |
| 150 | Returns: |
| 151 | The shape of the sharded version of the Tensor. |
| 152 | |
| 153 | Raises: |
| 154 | ValueError: If shard_index is None when shards are of different |
| 155 | shapes; or shard_index is not None and |
| 156 | !(0<=shard_index<number_of_shards); or shape does not have at |
| 157 | least self.shard_dimension+1 dimensions; or the value of |
| 158 | shape's shard dimension is not a multiple of |
| 159 | self.number_of_shards |
| 160 | """ |
| 161 | if self._shard_dimension is None or self._number_of_shards is None: |
| 162 | # Don't raise an error if the config is unset. |
| 163 | return None |
| 164 | if shard_index is not None: |
| 165 | if shard_index < 0 or shard_index >= self.number_of_shards: |
| 166 | raise ValueError("shard_index %d, but must be in [0,%d)." % |
| 167 | (shard_index, self._number_of_shards)) |
| 168 | shape = tensor_shape.as_shape(shape) |
| 169 | if self._number_of_shards == 1: |
| 170 | # Don't do anything when there's only one shard. |
| 171 | return shape |
| 172 | ndims = shape.ndims |
| 173 | if ndims is None: |
| 174 | raise ValueError("shape must be a specified shape not Unknown") |
| 175 | if ndims <= self._shard_dimension: |
| 176 | raise ValueError("shape %s does not contain shard_dimension %d" % |
| 177 | (shape.as_list(), self._shard_dimension)) |
| 178 | dims = shape.as_list() |
| 179 | if dims[self._shard_dimension] is None: |
| 180 | raise ValueError("shape %s must have a fixed size for dimension %d " |
| 181 | "that is known at graph construction time." % |
| 182 | (shape.as_list(), self._shard_dimension)) |
| 183 | if (dims[self._shard_dimension] % self._number_of_shards) != 0: |
| 184 | raise ValueError("shape %s cannot be sharded %d ways along dimension %d" % |
| 185 | (shape.as_list(), self._number_of_shards, |
| 186 | self._shard_dimension)) |
| 187 | dims[self._shard_dimension] /= self._number_of_shards |
| 188 | return tensor_shape.as_shape(dims) |
| 189 | |
| 190 | def _unshard_shape(self, shape): |
| 191 | """Return the unsharded shape that would generate a given sharded shape. |