Grouped allreduce, (across the given dimensions).

    Args:
      x: a LaidOutTensor
      mesh_axes: a list of integers - the mesh dimensions to be reduced
      reduction_fn_string: "SUM" or "MAX"
    Returns:
      a LaidOutTensor
183256
def allreduce(self, x, mesh_axes, reduction_fn_string):
    
    return self._collective_with_groups(
        x, mesh_axes, functools.partial(
            allreduce_ring, reduction_fn_string=reduction_fn_string))
1194304
def max(self, axis=None, dtype=None, out=None, keepdims=False):
        
        return self.elem.__array_ufunc__(
            np.maximum, 'reduce', self.elem,
            axis=axis, dtype=dtype, out=(out,), keepdims=keepdims)
807463
def laid_out_pcoord(self, mesh_axis):
    
    divisor = list_product(self.shape.to_integer_list[mesh_axis + 1:])
    modulus = self.shape[mesh_axis].size
    def my_fn(pnum):
      return (pnum // divisor) % modulus
    return self.slicewise(my_fn, self.laid_out_pnum())
709128
def row_coordinates(self, X):
        
        utils.validation.check_is_fitted(self, 's_')

        # Extract index
        index = X.index if isinstance(X, pd.DataFrame) else None

        # Copy data
        if self.copy:
            X = np.copy(X)

        # Scale data
        if hasattr(self, 'scaler_'):
            X = self.scaler_.transform(X)

        return pd.DataFrame(data=X.dot(self.V_.T), index=index)
218888
def _row_sparse_pull(self, parameter, out, row_id, full_idx=False):
        
        # initialize kv and params if not already
        if not self._kv_initialized:
            self._init_kvstore()
        if self._params_to_init:
            self._init_params()
        idx = self._param2idx[parameter.name]
        if full_idx and 'dist' not in self._kvstore.type:
            assert row_id.size == out.shape[0]
            self._kvstore.pull(idx, out=out, priority=-idx, ignore_sparse=False)
        else:
            self._kvstore.row_sparse_pull(idx, out=out, row_ids=row_id, priority=-idx)
1281477
def all_but_axis(i, axis, num_axes):
    
    the_slice  = ()
    for j in range(num_axes):
        if j == axis:
            the_slice = the_slice + (i,)
        else:
            the_slice = the_slice + (slice(None),)
    return the_slice
736417
def _prepare_axes(node, sort_key):
    
    links = node.links
    o_links = node._overlapping_links
    overlap = {ax2 for ax in links for ax2 in o_links.get(ax, [])}
    axes = []
    for axis in sorted(links.keys(), key=sort_key):
        if axis in overlap: continue
        tgt = links[axis]
        if axis in o_links:
            s, e = axis[0], axis[-1]
            axis = '%s%s%s' % (
                s, '&'.join(a[1:-1] for a in [axis] + o_links[axis]), e
            )
        axes.append((axis, tgt))
    return axes
17106
def combine_last_two_dimensions(x):
  
  x_shape = common_layers.shape_list(x)
  a, b = x_shape[-2:]
  return tf.reshape(x, x_shape[:-2] + [a * b])
806747
def _needs_reindex_multi(self, axes, method, level):
        """"""
        return ((com.count_not_none(*axes.values()) == self._AXIS_LEN) and
                method is None and level is None and not self._is_mixed_type)
169985
def import_tf_tensor(self, x, tf_x):
    
    return self.LaidOutTensor(self.make_slices(tf_x, x.shape))
344410
def _move_tuple_axes_first(array, axis):
    

    # Figure out how many axes we are operating over
    naxis = len(axis)

    # Add remaining axes to the axis tuple
    axis += tuple(i for i in range(array.ndim) if i not in axis)

    # The new position of each axis is just in order
    destination = tuple(range(array.ndim))

    # Reorder the array so that the axes being operated on are at the beginning
    array_new = np.moveaxis(array, axis, destination)

    # Figure out the size of the product of the dimensions being operated on
    first = np.prod(array_new.shape[:naxis])

    # Collapse the dimensions being operated on into a single dimension so that
    # we can then use axis=0 with the bottleneck functions
    array_new = array_new.reshape((first,) + array_new.shape[naxis:])

    return array_new
871911
def _merge_multi_context(outputs, major_axis):
    
    rets = []
    for tensors, axis in zip(outputs, major_axis):
        if axis >= 0:
            # pylint: disable=no-member,protected-access
            if len(tensors) == 1:
                rets.append(tensors[0])
            else:
                # Concatenate if necessary
                rets.append(nd.concat(*[tensor.as_in_context(tensors[0].context)
                                        for tensor in tensors],
                                      dim=axis))
            # pylint: enable=no-member,protected-access
        else:
            # negative axis means the there is no batch_size axis, and all the
            # results should be the same on each device. We simply take the
            # first one, without checking they are actually the same
            rets.append(tensors[0])
    return rets
286211
def table(self,x,*args,**kwargs):
        
        reduce_func = kwargs.get('reduce',lambda x: x)
        data = []
        for i,d in enumerate(self.generate(*args)):
            if len(d) > 0:
                data.append(reduce_func(d[x]))
            else:
                data.append(0)
        try:
            return np.reshape(data,[len(self.space(a)) for a in args])
        except:
            return np.reshape(data,[-1]+[len(self.space(a)) for a in args])
641204
def multigrid(bounds, points_count):
    
    if len(bounds)==1:
        return np.linspace(bounds[0][0], bounds[0][1], points_count).reshape(points_count, 1)
    x_grid_rows = np.meshgrid(*[np.linspace(b[0], b[1], points_count) for b in bounds])
    x_grid_columns = np.vstack([x.flatten(order='F') for x in x_grid_rows]).T
    return x_grid_columns
126468
def _easy_facetgrid(data, plotfunc, kind, x=None, y=None, row=None,
                    col=None, col_wrap=None, sharex=True, sharey=True,
                    aspect=None, size=None, subplot_kws=None, **kwargs):
    
    ax = kwargs.pop('ax', None)
    figsize = kwargs.pop('figsize', None)
    if ax is not None:
        raise ValueError("Can't use axes when making faceted plots.")
    if aspect is None:
        aspect = 1
    if size is None:
        size = 3
    elif figsize is not None:
        raise ValueError('cannot provide both `figsize` and `size` arguments')

    g = FacetGrid(data=data, col=col, row=row, col_wrap=col_wrap,
                  sharex=sharex, sharey=sharey, figsize=figsize,
                  aspect=aspect, size=size, subplot_kws=subplot_kws)

    if kind == 'line':
        return g.map_dataarray_line(plotfunc, x, y, **kwargs)

    if kind == 'dataarray':
        return g.map_dataarray(plotfunc, x, y, **kwargs)
402780
private int chooseMaximalExtendedSplitAxis(List<? extends SpatialComparable> objects) {
    // maximum and minimum value for the extension
    int dimension = objects.get(0).getDimensionality();
    double[] maxExtension = new double[dimension];
    double[] minExtension = new double[dimension];
    Arrays.fill(minExtension, Double.MAX_VALUE);

    // compute min and max value in each dimension
    for (SpatialComparable object : objects) {
      for (int d = 0; d < dimension; d++) {
        double min, max;
        min = object.getMin(d);
        max = object.getMax(d);

        if (maxExtension[d] < max) {
          maxExtension[d] = max;
        }

        if (minExtension[d] > min) {
          minExtension[d] = min;
        }
      }
    }

    // set split axis to dim with maximal extension
    int splitAxis = -1;
    double max = 0;
    for (int d = 0; d < dimension; d++) {
      double currentExtension = maxExtension[d] - minExtension[d];
      if (max < currentExtension) {
        max = currentExtension;
        splitAxis = d;
      }
    }
    return splitAxis;
  }
575914
def reduce_max(x,
               disable_positional_args=None,
               output_shape=None,
               reduced_dim=None,
               name=None):
  
  output_shape = convert_to_shape(output_shape)
  reduced_dim = convert_to_dimension(reduced_dim)
  assert disable_positional_args is None
  output_shape = _reduction_output_shape(x, output_shape, reduced_dim)
  if output_shape is None:
    output_shape = Shape([])
  if output_shape == x.shape:
    return x
  return ReduceOperation(
      x, output_shape, "MAX", name=name or "reduce_max").outputs[0]