(x1, x2)
| 196 | |
| 197 | def empirical_ntk_ntk_vps(func, params, x1, x2, compute='full'): |
| 198 | def get_ntk(x1, x2): |
| 199 | def func_x1(params): |
| 200 | return func(params, x1) |
| 201 | |
| 202 | def func_x2(params): |
| 203 | return func(params, x2) |
| 204 | |
| 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) |
| 217 | return vmap(get_ntk_slice)(basis) |
| 218 | |
| 219 | # ``get_ntk(x1, x2)`` computes the NTK for a single data point x1, x2 |
| 220 | # Since the x1, x2 inputs to ``empirical_ntk_ntk_vps`` are batched, |
nothing calls this directly
no outgoing calls
no test coverage detected