Computes and stacks jacobians of `output[i,...]` w.r.t. `input[i,...]`. e.g. x = tf.constant([[1, 2], [3, 4]], dtype=tf.float32) y = x * x jacobian = batch_jacobian(y, x) # => [[[2, 0], [0, 4]], [[6, 0], [0, 8]]] Args: output: A tensor with shape [b, y1, ..., y_n]. `output[i,..
(output, inp, use_pfor=True, parallel_iterations=None)
| 81 | |
| 82 | |
| 83 | def batch_jacobian(output, inp, use_pfor=True, parallel_iterations=None): |
| 84 | """Computes and stacks jacobians of `output[i,...]` w.r.t. `input[i,...]`. |
| 85 | |
| 86 | e.g. |
| 87 | x = tf.constant([[1, 2], [3, 4]], dtype=tf.float32) |
| 88 | y = x * x |
| 89 | jacobian = batch_jacobian(y, x) |
| 90 | # => [[[2, 0], [0, 4]], [[6, 0], [0, 8]]] |
| 91 | |
| 92 | Args: |
| 93 | output: A tensor with shape [b, y1, ..., y_n]. `output[i,...]` should |
| 94 | only depend on `inp[i,...]`. |
| 95 | inp: A tensor with shape [b, x1, ..., x_m] |
| 96 | use_pfor: If true, uses pfor for computing the Jacobian. Else uses a |
| 97 | tf.while_loop. |
| 98 | parallel_iterations: A knob to control how many iterations and dispatched in |
| 99 | parallel. This knob can be used to control the total memory usage. |
| 100 | |
| 101 | Returns: |
| 102 | A tensor `t` with shape [b, y_1, ..., y_n, x1, ..., x_m] where `t[i, ...]` |
| 103 | is the jacobian of `output[i, ...]` w.r.t. `inp[i, ...]`, i.e. stacked |
| 104 | per-example jacobians. |
| 105 | |
| 106 | Raises: |
| 107 | ValueError: if first dimension of `output` and `inp` do not match. |
| 108 | """ |
| 109 | output_shape = output.shape |
| 110 | if not output_shape[0].is_compatible_with(inp.shape[0]): |
| 111 | raise ValueError("Need first dimension of output shape (%s) and inp shape " |
| 112 | "(%s) to match." % (output.shape, inp.shape)) |
| 113 | if output_shape.is_fully_defined(): |
| 114 | batch_size = int(output_shape[0]) |
| 115 | output_row_size = output_shape.num_elements() // batch_size |
| 116 | else: |
| 117 | output_shape = array_ops.shape(output) |
| 118 | batch_size = output_shape[0] |
| 119 | output_row_size = array_ops.size(output) // batch_size |
| 120 | inp_shape = array_ops.shape(inp) |
| 121 | # Flatten output to 2-D. |
| 122 | with ops.control_dependencies( |
| 123 | [check_ops.assert_equal(batch_size, inp_shape[0])]): |
| 124 | output = array_ops.reshape(output, [batch_size, output_row_size]) |
| 125 | |
| 126 | def loop_fn(i): |
| 127 | y = array_ops.gather(output, i, axis=1) |
| 128 | return gradient_ops.gradients(y, inp)[0] |
| 129 | |
| 130 | if use_pfor: |
| 131 | pfor_output = control_flow_ops.pfor(loop_fn, output_row_size, |
| 132 | parallel_iterations=parallel_iterations) |
| 133 | else: |
| 134 | pfor_output = control_flow_ops.for_loop( |
| 135 | loop_fn, output.dtype, |
| 136 | output_row_size, |
| 137 | parallel_iterations=parallel_iterations) |
| 138 | if pfor_output is None: |
| 139 | return None |
| 140 | pfor_output = array_ops.reshape(pfor_output, |
nothing calls this directly
no test coverage detected