Adds tf options to numpy scan ops.
(func, x, axis, exclusive, reverse)
| 42 | |
| 43 | |
| 44 | def handle_options(func, x, axis, exclusive, reverse): |
| 45 | """Adds tf options to numpy scan ops.""" |
| 46 | length = len(x.shape) |
| 47 | if axis < 0: |
| 48 | axis = length + axis |
| 49 | |
| 50 | if reverse: |
| 51 | x = numpy_reverse(x, axis) |
| 52 | |
| 53 | if exclusive: |
| 54 | ix_head = [slice(0, 1) if i == axis else slice(None) for i in range(length)] |
| 55 | ix_init = [ |
| 56 | slice(0, -1) if i == axis else slice(None) for i in range(length) |
| 57 | ] |
| 58 | if func == np.cumsum: |
| 59 | init = np.zeros_like(x[ix_head]) |
| 60 | elif func == np.cumprod: |
| 61 | init = np.ones_like(x[ix_head]) |
| 62 | else: |
| 63 | raise ValueError("Unknown scan function.") |
| 64 | x = np.concatenate([init, func(x[ix_init], axis)], axis=axis) |
| 65 | else: |
| 66 | x = func(x, axis=axis) |
| 67 | |
| 68 | if reverse: |
| 69 | x = numpy_reverse(x, axis) |
| 70 | return x |
| 71 | |
| 72 | |
| 73 | class CumsumTest(test.TestCase): |
no test coverage detected