JAX
关注
分享
扫一扫
文章平均质量分 95
张博208
知识搬运工
展开
-
How does vmap() function work?
According to JAX docsfile on vmap, jax.vmap(function, in_axes=0, out_axes=0) returns a function which maps the function one specified over using in_axes and stack them together using out_axes原创 2024-06-04 11:44:32 · 676 阅读 · 0 评论 -
JAX VMAP Simplified: An Easy Introduction for Beginners
WMAP转载 2024-06-04 10:03:36 · 106 阅读 · 0 评论