Open tingwl0122 opened 3 months ago
For most of the algorithms we implemented, we use some specific dimension manipulation to make vmap work on batched gradient computation. We will fix this for all the algorithms afterward. (originally mentioned in #83 )
vmap
For most of the algorithms we implemented, we use some specific dimension manipulation to make
vmap
work on batched gradient computation. We will fix this for all the algorithms afterward. (originally mentioned in #83 )