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