[functorch] Update neural_tangent_kernels.ipynb (pytorch/functorch#788)
Fix a small bug
```python3
if compute == 'full':
return result
if compute == 'trace':
return torch.einsum('NMKK->NM') # should be torch.einsum('NMKK->NM', result)
if compute == 'diagonal':
return torch.einsum('NMKK->NMK') # should be torch.einsum('NMKK->NMK', result)
```
