(i)
| 51 | output_dtypes = [] |
| 52 | |
| 53 | def loop_fn(i): |
| 54 | with g: |
| 55 | x1 = array_ops.gather(x, i) |
| 56 | y1 = op(x1) |
| 57 | outputs = [op(x), y1] |
| 58 | if y1.dtype == dtypes.float32: |
| 59 | loss = math_ops.reduce_sum(y1 * y1) |
| 60 | else: |
| 61 | loss = None |
| 62 | if loss is not None: |
| 63 | grad = g.gradient(loss, x1) |
| 64 | if grad is not None: |
| 65 | outputs.append(grad) |
| 66 | del output_dtypes[:] |
| 67 | output_dtypes.extend([t.dtype for t in outputs]) |
| 68 | return outputs |
| 69 | |
| 70 | # pylint: enable=cell-var-from-loop |
| 71 |