hack in keepdims behavior into a function taking an axis
(func)
| 27 | |
| 28 | |
| 29 | def _add_keepdims(func): |
| 30 | """ hack in keepdims behavior into a function taking an axis """ |
| 31 | @functools.wraps(func) |
| 32 | def wrapped(a, axis, **kwargs): |
| 33 | res = func(a, axis=axis, **kwargs) |
| 34 | if axis is None: |
| 35 | axis = 0 # res is now a scalar, so we can insert this anywhere |
| 36 | return np.expand_dims(res, axis=axis) |
| 37 | return wrapped |
| 38 | |
| 39 | |
| 40 | class TestTakeAlongAxis: |
no outgoing calls
no test coverage detected
searching dependent graphs…