<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Manipulating axes
Parallelizing axes across multiple dimensions is a key concept for the distributed Bolt array. Bolt is designed so that you don't need to think about this most of the time, and everything is handled automatically. But it's worth understanding how things are structured under the hood, especially for more advanced usage.
Which axes are parallelized?
They core idea is to use a collection of key-value pairs where the keys and values separately represent axes of a single array. The keys are tuples that encode the indices of the "parallelized" axes, and the values are (smaller) ndarrays representing the "localized" axes. We can choose which axes to parallelize over during construction, and by convention, the key axes come before the value axes. The default is to parallelize over the first axis
Step1:  We can see the overall shape
Step2:  And we can also inspect the shapes of the keys and the values
Step3:  Because the key axes come before the value axes, we represent which axes are parallelized with a single number, the split, which in this case is 1
Step4:  For those who know Spark
Step5:  If we instead parallelize over the first two axes
Step6:  The overall shape is the same, but the shapes of the keys and values have changed
Step7:  Again, for those who know Spark, we can look at the actual records. We should now have length 2 tuples as the keys, and arrays of shape (4,) as the values.
Step8:  Transposing
A key operation on distributed arrays is transposing, which lets us apply parallelized operations over different axes
Due to the structure described above, transposes that involve either keys or values are simple because they just involve maps over the keys or the values. That's how the following operations are implemented (though this is handled for you under the hood)
Step9:  The more complex case is when the axes change between keys and values. This is handled through an efficient swap method that we developed. We chunk the values into subarrays of moderate size, group them, and reorganize them. The swap operation is more general than just a transpose, but it's what we use to perform transposes. In terms of performance, we've used this to efficiently transpose very large (TB+) arrays in minutes on moderately sized clusters.
The transpose itself is lazy, and shapes are computed automatically, so nothing happens until you request an output (usually by calling sum or toarray). Note the small difference in execution time between the following two operations; only the latter requires execution.
Step10:  Reshaping
Another key operation is reshaping. Like transpose, many reshapings can be performed independently on either the key or the value axes, and this is handled automatically given an arbtirary reshape.
Step11:  We have not yet solved the more general case (where the reshaping spans multiple axes).
Mapping over axes
As shown in the basic usage tutorial, we can apply a function over our distributed arrays in parallel using map. We can apply our function over any axes, but depending on the axes we choose, the implementation differs. If we map over the axes represented in the values, then it is simply a single map operation. If we map over another set of axes, it will require a swap to move those axes
For example, this does not incur a swap
Step12:  but this does
<ASSISTANT_TASK:>
Python Code:
from bolt import ones
a = ones((2, 3, 4), sc)
Explanation: Manipulating axes
Parallelizing axes across multiple dimensions is a key concept for the distributed Bolt array. Bolt is designed so that you don't need to think about this most of the time, and everything is handled automatically. But it's worth understanding how things are structured under the hood, especially for more advanced usage.
Which axes are parallelized?
They core idea is to use a collection of key-value pairs where the keys and values separately represent axes of a single array. The keys are tuples that encode the indices of the "parallelized" axes, and the values are (smaller) ndarrays representing the "localized" axes. We can choose which axes to parallelize over during construction, and by convention, the key axes come before the value axes. The default is to parallelize over the first axis
End of explanation
a.shape
Explanation: We can see the overall shape
End of explanation
a.keys.shape
a.values.shape
Explanation: And we can also inspect the shapes of the keys and the values
End of explanation
a.split
Explanation: Because the key axes come before the value axes, we represent which axes are parallelized with a single number, the split, which in this case is 1
End of explanation
a.tordd().keys().collect()
a.tordd().values().first()
Explanation: For those who know Spark: we are using an RDD where each record is a (tuple, ndarray) pair, where in this case the tuple is length 1, and the arrays have shape (3,4). When it's useful to look at the actual records, we can get them with the tordd method.
End of explanation
a = ones((2, 3, 4), sc, axis=(0, 1))
Explanation: If we instead parallelize over the first two axes
End of explanation
a.shape
a.keys.shape
a.values.shape
Explanation: The overall shape is the same, but the shapes of the keys and values have changed
End of explanation
a.tordd().keys().collect()
a.tordd().values().first()
Explanation: Again, for those who know Spark, we can look at the actual records. We should now have length 2 tuples as the keys, and arrays of shape (4,) as the values.
End of explanation
a = ones((2, 3, 4, 5), sc, axis=(0, 1))
a.transpose(0, 1, 3, 2).shape
a.transpose(1, 0, 2, 3).shape
Explanation: Transposing
A key operation on distributed arrays is transposing, which lets us apply parallelized operations over different axes
Due to the structure described above, transposes that involve either keys or values are simple because they just involve maps over the keys or the values. That's how the following operations are implemented (though this is handled for you under the hood)
End of explanation
%%time
a.transpose(3, 2, 1, 0).shape
%%time
a.transpose(3, 2, 1, 0).toarray().shape
Explanation: The more complex case is when the axes change between keys and values. This is handled through an efficient swap method that we developed. We chunk the values into subarrays of moderate size, group them, and reorganize them. The swap operation is more general than just a transpose, but it's what we use to perform transposes. In terms of performance, we've used this to efficiently transpose very large (TB+) arrays in minutes on moderately sized clusters.
The transpose itself is lazy, and shapes are computed automatically, so nothing happens until you request an output (usually by calling sum or toarray). Note the small difference in execution time between the following two operations; only the latter requires execution.
End of explanation
a.shape
a.reshape(3, 2, 4, 5).shape
a.reshape(2, 3, 5, 4).shape
Explanation: Reshaping
Another key operation is reshaping. Like transpose, many reshapings can be performed independently on either the key or the value axes, and this is handled automatically given an arbtirary reshape.
End of explanation
a = ones((2, 3, 4, 5), sc, axis=(0, 1))
a.map(lambda x: x * 2, axis=(0, 1)).toarray().shape
Explanation: We have not yet solved the more general case (where the reshaping spans multiple axes).
Mapping over axes
As shown in the basic usage tutorial, we can apply a function over our distributed arrays in parallel using map. We can apply our function over any axes, but depending on the axes we choose, the implementation differs. If we map over the axes represented in the values, then it is simply a single map operation. If we map over another set of axes, it will require a swap to move those axes
For example, this does not incur a swap
End of explanation
a.map(lambda x: x * 2, axis=(0,)).toarray().shape
Explanation: but this does
End of explanation
<END_TASK>