OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-07 12:21:27 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

[functorch] Update neural_tangent_kernels.ipynb (pytorch/functorch#788)

Browse Source 
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)
```
This commit is contained in:
Local State 2022-05-09 23:25:55 -04:00 committed by Jon Janzen
parent 126fd93c21
commit ff9558a2ea
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
functorch/notebooks/neural_tangent_kernels.ipynb
View File

@ -288,9 +288,9 @@
" if compute == 'full':\n",
" return result\n",
" if compute == 'trace':\n",
" return torch.einsum('NMKK->NM')\n",
" return torch.einsum('NMKK->NM', result)\n",
" if compute == 'diagonal':\n",
" return torch.einsum('NMKK->NMK')"
" return torch.einsum('NMKK->NMK', result)"
]
},
{