0

Say I have two numpy arrays, for example

import numpy as np

A = np.arange(5*3*3*2).reshape(5, 3, 3, 2)
B = np.arange(3*3).reshape(3, 3)

If I want to add A and B across a shared axis, I would just do

C = A + B[None, :, :, None]
# C has shape (5, 3, 3, 2) which is what I want

I want to write a write function that generalizes this kind of summation but am not how to get started. It would look something like

def mask(M, Mshape, out_shape):
    # not sure what to put here
    pass

def add_tensors(A, B, Ashape, Bshape, out_shape):
    # Here I mask A, B so that it has shape out_shape
    A = mask(A, Aaxis, out_shape)
    B = mask(B, Baxis, out_shape)
    return A + B

Any suggestions? Is it possible to make this a ufunc?

user3166083
  • 147
  • 8
  • Seems to me that `numpy's` broadcasting rules are as general as you can get. They have sorted out the ambiguity issues. If you don't like using the `None` you can also use `reshape` – hpaulj May 09 '20 at 21:21
  • I actually like using the ```None``` but am not sure how to apply this in a very general way (adding two tensors across a given axis). In the most extreme case, would ```reshape``` cause issues in performance due to memory allocation or something similar? – user3166083 May 09 '20 at 22:20
  • `reshape` returns a `view` where possible; just as basic indexing (with slice and None). – hpaulj May 09 '20 at 23:35

1 Answers1

0
In [447]: A = np.arange(5*3*3*2).reshape(5, 3, 3, 2) 
     ...: B = np.arange(3*3).reshape(3, 3)

These are all equivalent:

In [448]: A + B[None,:, :, None];                                                                      
In [449]: A + B[:, :, None];         # initial None is automatic

build the index tuple from list:

In [454]: tup = [slice(None)]*3; tup[-1] = None; tup = tuple(tup)                                      
In [455]: tup                                                                                          
Out[455]: (slice(None, None, None), slice(None, None, None), None)
In [456]: A + B[tup];

or the equivalent shape:

In [457]: sh = B.shape + (1,)                                                                          
In [458]: sh                                                                                           
Out[458]: (3, 3, 1)
In [459]: A + B.reshape(sh);

expand_dims also uses a parameterized reshape:

In [462]: np.expand_dims(B,2).shape                                                                    
Out[462]: (3, 3, 1)
In [463]: A+np.expand_dims(B,2);
hpaulj
  • 221,503
  • 14
  • 230
  • 353