MPI alltoall operation.

  Implementation of alltoall using pointwise communication.

  Args:
    xs: a list of n tf.Tensors
    devices: a list of n strings
    split_axis: an integer
    concat_axis: an integer

  Returns:
    a list of n Tensors
545197
def alltoall_pointtwise(xs, devices, split_axis, concat_axis):
  
  n = len(xs)
  if n == 1:
    return xs
  # [target, source]
  parts = mtf.transpose_list_of_lists(
      mtf.parallel(devices, tf.split, xs, [n] * n, axis=[split_axis] * n))
  return mtf.parallel(devices, tf.concat, parts, axis=[concat_axis] * n)
459547
def split_model_idx(model:nn.Module, idxs:Collection[int])->ModuleList:
    ""
    layers = flatten_model(model)
    if idxs[0] != 0: idxs = [0] + idxs
    if idxs[-1] != len(layers): idxs.append(len(layers))
    return [nn.Sequential(*layers[i:j]) for i,j in zip(idxs[:-1],idxs[1:])]
168926
def tensor(x:Any, *rest)->Tensor:
    ""
    if len(rest): x = (x,)+rest
    # XXX: Pytorch bug in dataloader using num_workers>0; TODO: create repro and report
    if is_listy(x) and len(x)==0: return tensor(0)
    res = torch.tensor(x) if is_listy(x) else as_tensor(x)
    if res.dtype is torch.int32:
        warn('Tensor is int32: upgrading to int64; for better performance use int64 input')
        return res.long()
    return res
137376
def Print(self, x, data, message, **kwargs):  # pylint: disable=invalid-name
    
    tf.logging.info("PlacementMeshImpl::Print")
    new_slices = x.tensor_list[:]
    with tf.device(self._devices[0]):
      new_slices[0] = tf.Print(
          new_slices[0], [t for d in data for t in d.tensor_list],
          message, **kwargs)
    return self.LaidOutTensor(new_slices)
1243132
def tensor_components_to_use(mrr, mtt, mpp, mrt, mrp, mtp):
    '''
    
    '''
    return np.array([[mrr, mrt, mrp], [mrt, mtt, mtp], [mrp, mtp, mpp]])
658077
def mergeSplitsOnInterfaces(root: LNode):
    
    for ch in root.children:
        if ch.children:
            mergeSplitsOnInterfaces(ch)

    ctx = MergeSplitsOnInterfacesCtx()
    for ch in root.children:
        srcPorts = None
        try:
            if ch.name == "CONCAT":
                p = single(ch.east, lambda x: True)
                e = single(p.outgoingEdges, lambda x: True)
                srcPorts = e.dsts
            elif ch.name == "SLICE":
                p = single(ch.west, lambda x: True)
                e = single(p.incomingEdges, lambda x: True)
                srcPorts = e.srcs
        except (DuplicitValueExc, NoValueExc):
            continue

        if srcPorts is not None:
            for srcPort in srcPorts:
                if isinstance(srcPort.parent, LPort):
                    # only for non primitive ports
                    rootPort = getRootIntfPort(srcPort)
                    ctx.register(rootPort, ch, e)

    # join them if it is possible
    for srcPort, splitsAndConcats in ctx.iterPortSplits():
        if len(splitsAndConcats) <= 1:
            continue

        name = "SPLIT" if srcPort.direction == PortType.OUTPUT else "CONCAT"
        newSplitNode = root.addNode(name)
        copyPort(srcPort, newSplitNode, True, "")
        n = splitsAndConcats[0][0]
        for i in range(max(len(n.west),
                           len(n.east))):
            copyPort(
                srcPort, newSplitNode,
                False, "[%d]" % i)

        reconnectPorts(root, srcPort, splitsAndConcats,
                       newSplitNode)
830958
def _multi_deref(tensors: List[tf.Tensor], index: tf.Tensor) -> List[tf.Tensor]:
    
    assert tensors
    assert tensors[0].shape[0] > 0
    return _deref_helper(lambda i: [tensor[i, ...] for tensor in tensors],
                         index,
                         0,
                         tensors[0].shape[0] - 1)
70887
def slice_begin(self, tensor_shape, pnum):
    
    tensor_layout = self.tensor_layout(tensor_shape)
    coordinates = pnum_to_processor_coordinates(self.shape, pnum)
    ret = []
    for dim_size, mesh_axis in zip(
        tensor_shape.to_integer_list, tensor_layout.tensor_axis_to_mesh_axis):
      if mesh_axis is None:
        ret.append(0)
      else:
        ret.append(
            dim_size // self.shape[mesh_axis].size * coordinates[mesh_axis])
    return ret
746258
def show_xys(self, xs, ys, imgsize:int=4, figsize:Optional[Tuple[int,int]]=None, **kwargs):
        ""
        rows = int(np.ceil(math.sqrt(len(xs))))
        axs = subplots(rows, rows, imgsize=imgsize, figsize=figsize)
        for x,y,ax in zip(xs, ys, axs.flatten()): x.show(ax=ax, y=y, **kwargs)
        for ax in axs.flatten()[len(xs):]: ax.axis('off')
        plt.tight_layout()
730629
def _ts_slice(self, Xt, y):
        '''  '''
        Ns = len(Xt)
        Xt_new = []
        for i in range(self.n_splits):
            for j in range(Ns):
                Njs = int(len(Xt[j]) / self.n_splits)
                Xt_new.append(Xt[j][(Njs * i):(Njs * (i + 1))])
        Xt_new = np.array(Xt_new)

        if len(np.atleast_1d(y[0])) == len(Xt[0]):
            # y is a time series
            y_new = []
            for i in range(self.n_splits):
                for j in range(Ns):
                    Njs = int(len(y[j]) / self.n_splits)
                    y_new.append(y[j][(Njs * i):(Njs * (i + 1))])
            y_new = np.array(y_new)
        else:
            # y is contextual to each series
            y_new = np.concatenate([y for i in range(self.n_splits)])

        return Xt_new, y_new
898410
def all(self, **kwargs):
        
        if self._is_transposed:
            # Pandas ignores on axis=1
            kwargs["bool_only"] = False
            kwargs["axis"] = kwargs.get("axis", 0) ^ 1
            return self.transpose().all(**kwargs)
        return self._process_all_any(lambda df, **kwargs: df.all(**kwargs), **kwargs)
986832
def All(a, axis, keep_dims):
    
    return np.all(a, axis=axis if not isinstance(axis, np.ndarray) else tuple(axis),
                  keepdims=keep_dims),
1005987
def _initialize_splittable_dimensions(self, mtf_graph):
    
    all_mtf_dimension_names = set()  # set(string)
    for mtf_operation in mtf_graph.operations:
      for mtf_tensor in mtf_operation.outputs:
        for mtf_dimension in mtf_tensor.shape.dims:
          if not re.match(r"_anonymous_\d*", mtf_dimension.name):
            all_mtf_dimension_names.add(mtf_dimension.name)

    unsplittable_mtf_dimension_names = set()  # set(string)
    for mtf_operation in mtf_graph.operations:
      unsplittable_mtf_dimension_names.update(mtf_operation.unsplittable_dims)

    return all_mtf_dimension_names - unsplittable_mtf_dimension_names
1365574
def tensor_vstack(tensor_list, pad=0):
    
    if len(tensor_list) == 1:
        return tensor_list[0][np.newaxis, :]

    ndim = len(tensor_list[0].shape)
    dimensions = [len(tensor_list)]  # first dim is batch size
    for dim in range(ndim):
        dimensions.append(max([tensor.shape[dim] for tensor in tensor_list]))

    dtype = tensor_list[0].dtype
    if pad == 0:
        all_tensor = np.zeros(tuple(dimensions), dtype=dtype)
    elif pad == 1:
        all_tensor = np.ones(tuple(dimensions), dtype=dtype)
    else:
        all_tensor = np.full(tuple(dimensions), pad, dtype=dtype)
    if ndim == 1:
        for ind, tensor in enumerate(tensor_list):
            all_tensor[ind, :tensor.shape[0]] = tensor
    elif ndim == 2:
        for ind, tensor in enumerate(tensor_list):
            all_tensor[ind, :tensor.shape[0], :tensor.shape[1]] = tensor
    elif ndim == 3:
        for ind, tensor in enumerate(tensor_list):
            all_tensor[ind, :tensor.shape[0], :tensor.shape[1], :tensor.shape[2]] = tensor
    else:
        raise Exception('Sorry, unimplemented.')
    return all_tensor
261178
def transform(self, X):
        

        extracted = []
        for columns, transformer in self.mapping:
            if transformer is not None:
                feature = transformer.transform(self._get_columns(X, columns))
            else:
                feature = self._get_columns(X, columns)

            if hasattr(feature, 'toarray'):
                # sparse arrays should be converted to regular arrays for hstack.
                feature = feature.toarray()

            if len(feature.shape) == 1:
                feature = np.array([feature]).T
            extracted.append(feature)

        # combine the feature outputs into one array.
        # at this point we lose track of which features
        # were created from which input columns, so it's
        # assumed that that doesn't matter to the model.
        return np.hstack(extracted)
630473
def nx_all_nodes_between(graph, source, target, data=False):
    
    import utool as ut
    if source is None:
        # assume there is a single source
        sources = list(ut.nx_source_nodes(graph))
        assert len(sources) == 1, (
            'specify source if there is not only one')
        source = sources[0]
    if target is None:
        # assume there is a single source
        sinks = list(ut.nx_sink_nodes(graph))
        assert len(sinks) == 1, (
            'specify sink if there is not only one')
        target = sinks[0]
    all_simple_paths = list(nx.all_simple_paths(graph, source, target))
    nodes = sorted(set.union(*map(set, all_simple_paths)))
    return nodes
686898
def export_to_tf_tensor(self, x):
    
    mesh_impl = self.mesh_impl(x)
    return mesh_impl.export_to_tf_tensor(
        x, self.tensors[x].to_laid_out_tensor())