Partitions or replicates the input tensor. The ops inside this function are placed on the host side. Args: tensor: The input tensor which will be partitioned or replicated. dims: A list of integer describes how to partition the input tensor. Returns: An iterator of `Tensor`s o
(tensor, dims)
| 38 | |
| 39 | |
| 40 | def partition_or_replicate_on_host(tensor, dims): |
| 41 | """Partitions or replicates the input tensor. |
| 42 | |
| 43 | The ops inside this function are placed on the host side. |
| 44 | |
| 45 | Args: |
| 46 | tensor: The input tensor which will be partitioned or replicated. |
| 47 | dims: A list of integer describes how to partition the input tensor. |
| 48 | |
| 49 | Returns: |
| 50 | An iterator of `Tensor`s or a list of partitioned tensors. |
| 51 | """ |
| 52 | if dims is None: |
| 53 | return itertools.repeat(tensor) |
| 54 | dims = np.array(dims) |
| 55 | output = [tensor] |
| 56 | shape_list = np.array(tensor.shape.as_list()) |
| 57 | quotients, remainders = np.divmod(shape_list, dims) |
| 58 | for axis, (quotient, remainder, dim, original_size) in enumerate( |
| 59 | zip(quotients, remainders, dims, shape_list)): |
| 60 | if dim <= 1: |
| 61 | continue |
| 62 | if remainder > 0: |
| 63 | # For each dimension, when it cannot be evenly partitioned, XLA assumes |
| 64 | # tensors are partitioned in a greedy manner by using |
| 65 | # ceil_ratio(size/dim) first. E.g. 2D tensor with shape (5, 14) and dims |
| 66 | # are (2, 4). Since 5 % 2 = 1 and 14 % 4 = 2, [5, 14] => |
| 67 | # [[(3, 4), (3, 4), (2, 4), (2, 2)], |
| 68 | # [(2, 4), (2, 4), (2, 4), (2, 2)]] |
| 69 | ceil_ratio = quotient + 1 |
| 70 | num_full_slots, left_over = np.divmod(original_size, ceil_ratio) |
| 71 | num_or_size_splits = [ceil_ratio] * num_full_slots + [left_over] |
| 72 | if len(num_or_size_splits) < dim: |
| 73 | num_or_size_splits += [0] * (dim - len(num_or_size_splits)) |
| 74 | new_output = [] |
| 75 | for x in output: |
| 76 | new_output.append( |
| 77 | array_ops.split( |
| 78 | x, num_or_size_splits=num_or_size_splits, axis=axis)) |
| 79 | output = new_output |
| 80 | else: |
| 81 | output = [array_ops.split(x, int(dim), axis=axis) for x in output] |
| 82 | output = nest.flatten(output) |
| 83 | return output |
| 84 | |
| 85 | |
| 86 | def _tag_sharding_attribute_for_dequeued_tensor(tensor, dims): |