hack in keepdims behavior into a function taking an axis
(func)
| 17 | |
| 18 | |
| 19 | def _add_keepdims(func): |
| 20 | """ hack in keepdims behavior into a function taking an axis """ |
| 21 | @functools.wraps(func) |
| 22 | def wrapped(a, axis, **kwargs): |
| 23 | res = func(a, axis=axis, **kwargs) |
| 24 | if axis is None: |
| 25 | axis = 0 # res is now a scalar, so we can insert this anywhere |
| 26 | return np.expand_dims(res, axis=axis) |
| 27 | return wrapped |
| 28 | |
| 29 | |
| 30 | class TestTakeAlongAxis: |
no outgoing calls
no test coverage detected