Compute the jacobian of `f` at `primals` using forward-mode autodiff.
(f, primals)
| 56 | |
| 57 | |
| 58 | def _jacfwd(f, primals): |
| 59 | """Compute the jacobian of `f` at `primals` using forward-mode autodiff.""" |
| 60 | jac_flat = [] |
| 61 | flat_primals = nest.flatten(primals) |
| 62 | tangent_mask = [array_ops.zeros_like(primal) for primal in flat_primals] |
| 63 | for primal_index, primal in enumerate(flat_primals): |
| 64 | primal_vector = array_ops.reshape(primal, [-1]) |
| 65 | primal_vector_length = array_ops.size(primal_vector) |
| 66 | jac_columns = [] |
| 67 | for element_index in math_ops.range(primal_vector_length): |
| 68 | mask = array_ops.one_hot(element_index, primal_vector_length) |
| 69 | tangent_mask[primal_index] = array_ops.reshape(mask, |
| 70 | array_ops.shape(primal)) |
| 71 | jac_columns.append( |
| 72 | nest.map_structure( |
| 73 | functools.partial(array_ops.reshape, shape=[-1]), |
| 74 | _jvp(f, primals, tangent_mask)[1])) |
| 75 | jac_flat.append(array_ops.stack(jac_columns, axis=1)) |
| 76 | tangent_mask[primal_index] = array_ops.zeros_like(primal) |
| 77 | return nest.pack_sequence_as(primals, jac_flat) |
| 78 | |
| 79 | |
| 80 | def _grad(f, argnums=0): |