import numpy as np


def compute_subtensor_shapes_unbalanced(local_tensor_structure, P_in, P_out=None):
    r"""Collectively assembles the shapes of all subtensors in a partition
        with no prior assumption on the distribution of elements.

        The output array has shape D+1, where D is the dimension of the tensor
        in question.  The first D dimensions of `shapes` correspond to the
        Cartesian decomposed workers and the last dimension corresponds to
        the shape of each subtensor.

        Parameters
        ----------
        local_tensor_structure : TensorStructure
            Structure of the local tensor, to be partially shared with P_out.
        P_in : MPIPartition
            Partition containing the local tensor.
        P_out : MPIPartition, optional
            Partition where an array of the local subtensor shapes will end up.

        Warning
        -------
        It is assumed that P_out has at least P_in's rank 0 in common.


        Returns
        -------
        Cartesian indexed NumPy array of all subtensors sizes on P_in.

    """

    shapes = None

    if P_in.active:

        shapes_size = list(P_in.shape) + [P_in.dim]
        shapes = np.zeros(shapes_size, dtype=np.int)

        # Generate a slice that isolates the worker's index into shapes
        # where it can store its shape
        sl = tuple([slice(i, i+1) for i in P_in.index] + [slice(None)])
        shapes[sl] = np.asarray(local_tensor_structure.shape)

        # Share everyone's shapes
        shapes = P_in.allreduce_data(shapes)

        # all max-reduce the shape

    if P_out is not None and P_out.active:
        shapes = P_out.broadcast_data(shapes, P_data=P_in)

    return shapes

import numpy as np


def compute_subtensor_shapes_unbalanced(local_tensor_structure, P_in, P_out=None):
    r"""Collectively assembles the shapes of all subtensors in a partition
        with no prior assumption on the distribution of elements.

        The output array has shape D+1, where D is the dimension of the tensor
        in question.  The first D dimensions of `shapes` correspond to the
        Cartesian decomposed workers and the last dimension corresponds to
        the shape of each subtensor.

        Parameters
        ----------
        local_tensor_structure : TensorStructure
            Structure of the local tensor, to be partially shared with P_out.
        P_in : MPIPartition
            Partition containing the local tensor.
        P_out : MPIPartition, optional
            Partition where an array of the local subtensor shapes will end up.

        Warning
        -------
        It is assumed that P_out has at least P_in's rank 0 in common.


        Returns
        -------
        Cartesian indexed NumPy array of all subtensors sizes on P_in.

    """

    shapes = None

    if P_in.active:

        shapes_size = list(P_in.shape) + [P_in.dim]
        shapes = np.zeros(shapes_size, dtype=np.int)

        # Generate a slice that isolates the worker's index into shapes
        # where it can store its shape
        sl = tuple([slice(i, i+1) for i in P_in.index] + [slice(None)])
        shapes[sl] = np.asarray(local_tensor_structure.shape)

        # Share everyone's shapes
        shapes = P_in.allreduce_data(shapes)

        # all max-reduce the shape

    if P_out is not None and P_out.active:
        shapes = P_out.broadcast_data(shapes, P_data=P_in)

    return shapes

Collectively assembles the shapes of all subtensors in a partition         with no prior assumption on the distribution of elements.          The output array has shape D+1, where D is the dimension of the tensor         in question.  The first D dimensions of `shapes` correspond to the         Cartesian decomposed workers and the last dimension corresponds to         the shape of each subtensor.          Parameters         ----------         local_tensor_structure : TensorStructure             Structure of the local tensor, to be partially shared with P_out.         P_in : MPIPartition             Partition containing the local tensor.         P_out : MPIPartition, optional             Partition where an array of the local subtensor shapes will end up.          Warning         -------         It is assumed that P_out has at least P_in's rank 0 in common.           Returns         -------         Cartesian indexed NumPy array of all subtensors sizes on P_in. # Generate a slice that isolates the worker's index into shapes # where it can store its shape # Share everyone's shapes # all max-reduce the shape