(vec)
| 205 | output, vjp_fn = vjp(func_x1, params) |
| 206 | |
| 207 | def get_ntk_slice(vec): |
| 208 | # This computes ``vec @ J(x2).T`` |
| 209 | # `vec` is some unit vector (a single slice of the Identity matrix) |
| 210 | vjps = vjp_fn(vec) |
| 211 | # This computes ``J(X1) @ vjps`` |
| 212 | _, jvps = jvp(func_x2, (params,), vjps) |
| 213 | return jvps |
| 214 | |
| 215 | # Here's our identity matrix |
| 216 | basis = torch.eye(output.numel(), dtype=output.dtype, device=output.device).view(output.numel(), -1) |
nothing calls this directly
no outgoing calls
no test coverage detected