Predict memory layout of ufunc output

Question

Using numpy ndarrays most of the time we need't worry our pretty little heads about memory layout because results do not depend on it.

Except when they do. Consider, for example, this slightly overengineered way of setting the diagonal of a 3x2 matrix

>>> a = np.zeros((3,2))
>>> a.reshape(2,3)[:,0] = 1
>>> a
array([[1., 0.],
       [0., 1.],
       [0., 0.]])

As long as we control the memory layout of a this is fine. But if we don't it is a bug and to make matters worse a nasty silent one:

>>> a = np.zeros((3,2),order='F')
>>> a.reshape(2,3)[:,0] = 1
>>> a
array([[0., 0.],
       [0., 0.],
       [0., 0.]])

This shall suffice to show that memory layout is not merely an implementation detail.

The first thing one might reasonably ask to get on top of array layout is What do new arrays look like? The factories empty,ones,zeros,identity etc. return C-contiguous layouts per default.

This rule does, however, not extend to every new array that was allocated by numpy. For example:

>>> a = np.arange(8).reshape(2,2,2).transpose(1,0,2)
>>> aa = a*a

The product aa is a new array allocated by ufunc np.multiply. Is it C-contiguous? No:

>>> aa.strides
(16, 32, 8)

My guess is that this is the result of an optimization that recognizes that this operation can be done on a flat linear array which would explain why the output has the same memory layout as the inputs.

In fact this can even be useful, unlike the following nonsense function. It shows a handy idiom to implement an axis parameter while still keeping indexing simple.

>>> def symmetrize_along_axis(a,axis=0):
...     aux = a.swapaxes(0,axis)
...     out = aux + aux[::-1]
...     return out.swapaxes(0,axis)

The slightly surprising but clearly desirable thing is that this produces contiguous output as long as input is contiguous.

>>> a = np.arange(8).reshape(2,2,2)
>>> symmetrize_along_axis(a,1).flags.contiguous
True

This shall suffice to show that knowing what layouts are returned by ufuncs can be quite useful. Hence my question:

Given the layouts of ufunc arguments are there any rules or guarantees regarding the layout of the output?

Answer 1

In a = np.zeros((3,2),order='F') case, a.reshape(2,3) creates a copy, not a view. That's why assignment fails, not the memory layout itself.

Look at same shape array:

In [123]: a = np.arange(6).reshape(3,2)                                                              
In [124]: a                                                                                          
Out[124]: 
array([[0, 1],
       [2, 3],
       [4, 5]])
In [125]: a.reshape(2,3)                                                                             
Out[125]: 
array([[0, 1, 2],
       [3, 4, 5]])
In [127]: a.reshape(2,3)[:,0]                                                                        
Out[127]: array([0, 3])

In [125] the values still flow in order C.

and an order F array:

In [128]: b = np.arange(6).reshape(3,2, order='F')                                                   
In [129]: b                                                                                          
Out[129]: 
array([[0, 3],                 # values flow in order F
       [1, 4],
       [2, 5]])
In [130]: b.reshape(2,3)                                                                             
Out[130]: 
array([[0, 3, 1],              # values are jumbled
       [4, 2, 5]])
In [131]: b.reshape(2,3)[:,0]                                                                        
Out[131]: array([0, 4])

If I keep order F in the shape:

In [132]: b.reshape(2,3, order='F')                                                                  
Out[132]: 
array([[0, 2, 4],               # values still flow in order F
       [1, 3, 5]])
In [133]: b.reshape(2,3, order='F')[:,0]                                                             
Out[133]: array([0, 1])

Confirm with assignment:

In [135]: a.reshape(2,3)[:,0]=10                                                                     
In [136]: a                                                                                          
Out[136]: 
array([[10,  1],
       [ 2, 10],
       [ 4,  5]])

not assignment:

In [137]: b.reshape(2,3)[:,0]=10                                                                     
In [138]: b                                                                                          
Out[138]: 
array([[0, 3],
       [1, 4],
       [2, 5]])

but here assignment works:

In [139]: b.reshape(2,3, order='F')[:,0]=10                                                          
In [140]: b                                                                                          
Out[140]: 
array([[10,  3],
       [10,  4],
       [ 2,  5]])

Or we can use order A to preserve order:

In [143]: b.reshape(2,3, order='A')[:,0]                                                             
Out[143]: array([10, 10])
In [144]: b.reshape(2,3, order='A')[:,0] = 20                                                        
In [145]: b                                                                                          
Out[145]: 
array([[20,  3],
       [20,  4],
       [ 2,  5]])

ufunc order

Suspecting that ufunc are (mostly) implemented with nditer (C version), I checked np.nditer` docs - order can be specified in several places. And the tutorial demonstrates order effect on the iteration.

I don't see order documented for ufunc, but, it is accepted by the kwargs.

In [171]: c = np.arange(8).reshape(2,2,2)                                                            
In [172]: d = c.transpose(1,0,2)                                                                     
In [173]: d.strides                                                                                  
Out[173]: (16, 32, 8)
In [174]: np.multiply(d,d,order='K').strides                                                         
Out[174]: (16, 32, 8)
In [175]: np.multiply(d,d,order='C').strides                                                         
Out[175]: (32, 16, 8)
In [176]: np.multiply(d,d,order='F').strides                                                         
Out[176]: (8, 16, 32)