Computes the numeric Jacobian for dy/dx. Computes the numeric Jacobian by slightly perturbing the inputs and measuring the differences on the output. Args: x: the tensor "x". x_shape: the dimensions of x as a tuple or an array of ints. x_data: a numpy array as the input data for
(x, x_shape, x_data, y, y_shape, delta,
extra_feed_dict)
| 133 | |
| 134 | |
| 135 | def _compute_numeric_jacobian(x, x_shape, x_data, y, y_shape, delta, |
| 136 | extra_feed_dict): |
| 137 | """Computes the numeric Jacobian for dy/dx. |
| 138 | |
| 139 | Computes the numeric Jacobian by slightly perturbing the inputs and |
| 140 | measuring the differences on the output. |
| 141 | |
| 142 | Args: |
| 143 | x: the tensor "x". |
| 144 | x_shape: the dimensions of x as a tuple or an array of ints. |
| 145 | x_data: a numpy array as the input data for x |
| 146 | y: the tensor "y". |
| 147 | y_shape: the dimensions of y as a tuple or an array of ints. |
| 148 | delta: the amount of perturbation we give to the input |
| 149 | extra_feed_dict: dict that allows fixing specified tensor values |
| 150 | during the jacobian calculation. |
| 151 | |
| 152 | Returns: |
| 153 | A 2-d numpy array representing the Jacobian for dy/dx. It has "x_size" rows |
| 154 | and "y_size" columns where "x_size" is the number of elements in x and |
| 155 | "y_size" is the number of elements in y. |
| 156 | """ |
| 157 | # bfloat16 doesn't have enough bits to represent high precision numbers such |
| 158 | # as delta. Convert to float32 here. Since numeric_jacobian is expected to |
| 159 | # be the groundtruth to compare against, it shouldn't lose any information. |
| 160 | if x.dtype == dtypes.bfloat16: |
| 161 | x = math_ops.cast(x, dtypes.float32) # TODO(wangpeng): Now that the new x |
| 162 | # is an output of the old x, isn't feeding to the new x a mistake? |
| 163 | if y.dtype == dtypes.bfloat16: |
| 164 | y = math_ops.cast(y, dtypes.float32) |
| 165 | if x_data.dtype == dtypes.bfloat16.as_numpy_dtype: |
| 166 | x_data = x_data.astype(np.float32) |
| 167 | |
| 168 | # To compute the jacobian, we treat x and y as one-dimensional vectors |
| 169 | x_size = _product(x_shape) * (2 if x.dtype.is_complex else 1) |
| 170 | y_size = _product(y_shape) * (2 if y.dtype.is_complex else 1) |
| 171 | x_dtype = x.dtype.real_dtype.as_numpy_dtype |
| 172 | y_dtype = y.dtype.real_dtype.as_numpy_dtype |
| 173 | |
| 174 | # Make sure we have the right types |
| 175 | x_data = np.asarray(x_data, dtype=x.dtype.as_numpy_dtype) |
| 176 | scale = np.asarray(2 * delta, dtype=y_dtype)[()] |
| 177 | |
| 178 | jacobian = np.zeros((x_size, y_size), dtype=x_dtype) |
| 179 | # For each of the entry of x, we slightly perturbs this by adding and |
| 180 | # subtracting a delta and then compute difference between the outputs. This |
| 181 | # will give us one row of the Jacobian matrix. |
| 182 | for row in range(x_size): |
| 183 | x_pos = x_data.copy() |
| 184 | x_neg = x_data.copy() |
| 185 | x_pos.ravel().view(x_dtype)[row] += delta |
| 186 | y_pos = y.eval(feed_dict=_extra_feeds(extra_feed_dict, {x: x_pos})) |
| 187 | x_neg.ravel().view(x_dtype)[row] -= delta |
| 188 | y_neg = y.eval(feed_dict=_extra_feeds(extra_feed_dict, {x: x_neg})) |
| 189 | diff = (y_pos - y_neg) / scale |
| 190 | jacobian[row, :] = diff.ravel().view(y_dtype) |
| 191 | |
| 192 | logging.vlog(1, "Numeric Jacobian =\n%s", jacobian) |
no test coverage detected