{"id":2551,"date":"2026-02-15T18:47:27","date_gmt":"2026-02-15T10:47:27","guid":{"rendered":"https:\/\/ai.ziyuanzz.online\/?page_id=2551"},"modified":"2026-03-28T09:28:29","modified_gmt":"2026-03-28T01:28:29","slug":"jax","status":"publish","type":"page","link":"https:\/\/ai.ziyuanzz.online\/index.php\/jax\/","title":{"rendered":"JAX"},"content":{"rendered":"
\n
\"\"<\/figure>\n<\/div>\n\n\n

JAX<\/h2>\n\n\n\n

Google\u63a8\u51fa\u7684\u7528\u4e8e\u53d8\u6362\u6570\u503c\u51fd\u6570\u7684\u673a\u5668\u5b66\u4e60\u6846\u67b6<\/p>\n\n\n\n

\n
\u8bbf\u95ee\u5b98\u7f51 ><\/a><\/div>\n<\/div>\n\n\n\n
<\/div>\n\n\n\n
\"\"<\/figure>\n\n\n\n
<\/div>\n\n\n\n

JAX\u662f\u4ec0\u4e48<\/h2>\n\n\n\n

JAX\u662fGoogle\u63a8\u51fa\u7684\u9ad8\u6027\u80fd\u6570\u503c\u8ba1\u7b97\u5e93\uff0c\u63d0\u4f9b\u7c7b\u4f3cNumPy\u7684API\uff0c\u652f\u6301GPU\/TPU\u52a0\u901f\u3001\u81ea\u52a8\u5fae\u5206\u3001\u5373\u65f6\u7f16\u8bd1\uff08JIT\uff09\u548c\u5411\u91cf\u5316\u7b49\u529f\u80fd\u3002JAX\u901a\u8fc7XLA\uff08\u52a0\u901f\u7ebf\u6027\u4ee3\u6570\uff09\u7f16\u8bd1\u5668\u4f18\u5316\u4ee3\u7801\uff0c\u663e\u8457\u63d0\u5347\u8fd0\u884c\u6548\u7387\uff0c\u5728\u5927\u89c4\u6a21\u6570\u636e\u5904\u7406\u548c\u673a\u5668\u5b66\u4e60\u4e2d\u8868\u73b0\u7a81\u51fa\u3002JAX\u652f\u6301\u81ea\u52a8\u5fae\u5206\uff0c\u80fd\u8f7b\u677e\u8ba1\u7b97\u51fd\u6570\u68af\u5ea6\uff0c\u9002\u7528\u4e8e\u4f18\u5316\u7b97\u6cd5\u3002JAX\u7684\u5f02\u6b65\u6267\u884c\u6a21\u5f0f\u548c\u4e0d\u53ef\u53d8\u6570\u7ec4\u8bbe\u8ba1\u4f7f\u5176\u5728\u6027\u80fd\u548c\u53ef\u9760\u6027\u4e0a\u4f18\u4e8e\u4f20\u7edfNumPy<\/a>\uff0c\u662f\u73b0\u4ee3\u79d1\u5b66\u8ba1\u7b97\u548c\u673a\u5668\u5b66\u4e60\u7814\u7a76\u4e2d\u7684\u91cd\u8981\u5de5\u5177\u3002<\/p>\n\n\n\n

\"JAX\"<\/a><\/figure>\n\n\n\n

JAX\u7684\u4e3b\u8981\u529f\u80fd<\/h2>\n\n\n\n
    \n
  • \u81ea\u52a8\u5fae\u5206<\/strong>\uff1a\u901a\u8fc7jax.grad<\/code>\u7b49\u51fd\u6570\u81ea\u52a8\u8ba1\u7b97\u51fd\u6570\u7684\u68af\u5ea6\uff0c\u652f\u6301\u9ad8\u9636\u5bfc\u6570\uff0c\u5e7f\u6cdb\u5e94\u7528\u5728\u673a\u5668\u5b66\u4e60\u4e2d\u7684\u6a21\u578b\u8bad\u7ec3\u3002<\/li>\n\n\n\n
  • \u5373\u65f6\u7f16\u8bd1\uff08JIT\uff09<\/strong>\uff1a\u7528jax.jit<\/code>\u5c06Python\u51fd\u6570\u7f16\u8bd1\u6210\u4f18\u5316\u540e\u7684\u673a\u5668\u4ee3\u7801\uff0c\u663e\u8457\u63d0\u5347\u8fd0\u884c\u6548\u7387\uff0c\u5728\u5927\u89c4\u6a21\u8ba1\u7b97\u4e2d\u6548\u679c\u663e\u8457\u3002<\/li>\n\n\n\n
  • \u5411\u91cf\u5316<\/strong>\uff1a\u901a\u8fc7jax.vmap<\/code>\u81ea\u52a8\u5c06\u51fd\u6570\u5411\u91cf\u5316\uff0c\u907f\u514d\u624b\u52a8\u5faa\u73af\uff0c\u63d0\u9ad8\u4ee3\u7801\u6548\u7387\u548c\u53ef\u8bfb\u6027\u3002<\/li>\n\n\n\n
  • \u5e76\u884c\u5316<\/strong>\uff1a\u7528jax.pmap<\/code>\u652f\u6301\u8de8\u591a\u4e2a\u8bbe\u5907\uff08\u5982GPU\u3001TPU\uff09\u7684\u5e76\u884c\u8ba1\u7b97\uff0c\u52a0\u901f\u5927\u89c4\u6a21\u4efb\u52a1\u5904\u7406\u3002<\/li>\n\n\n\n
  • \u786c\u4ef6\u52a0\u901f<\/strong>\uff1a\u652f\u6301\u5728CPU\u3001GPU\u548cTPU\u4e0a\u8fd0\u884c\u4ee3\u7801\uff0c\u5145\u5206\u5229\u7528\u786c\u4ef6\u7684\u5e76\u884c\u8ba1\u7b97\u80fd\u529b\u3002<\/li>\n\n\n\n
  • \u7a0b\u5e8f\u53d8\u6362<\/strong>\uff1a\u63d0\u4f9b\u4e30\u5bcc\u7684\u7a0b\u5e8f\u53d8\u6362\u5de5\u5177\uff0c\u5982jax.lax<\/code>\uff0c\u7528\u5728\u6784\u5efa\u66f4\u590d\u6742\u7684\u7a0b\u5e8f\u903b\u8f91\uff0c\u63d0\u5347\u4ee3\u7801\u7075\u6d3b\u6027\u548c\u6269\u5c55\u6027\u3002<\/li>\n<\/ul>\n\n\n\n

    \u5982\u4f55\u4f7f\u7528JAX<\/h2>\n\n\n\n
      \n
    • \u73af\u5883\u914d\u7f6e\u4e0e\u5b89\u88c5<\/strong>\uff1a\n
        \n
      • \u521b\u5efaPython\u73af\u5883<\/strong>\uff1a\u7528conda\u521b\u5efa\u4e00\u4e2a\u4e13\u7528\u7684Python\u73af\u5883\u3002<\/li>\n<\/ul>\n<\/li>\n<\/ul>\n\n\n\n
        conda create <span class<\/strong>=\"token parameter variable\">-n<\/span> jax_test <span class<\/strong>=\"token assign-left variable\">python<\/span><span class<\/strong>=\"token operator\">=<\/span><span class<\/strong>=\"token number\">3.13<\/span> <span class<\/strong>=\"token parameter variable\">-y<\/span>conda activate jax_test<\/pre>\n\n\n\n
          \n
        • \n
            \n
          • \u5b89\u88c5JAX\u5e93<\/strong>\uff1a\u6839\u636e\u786c\u4ef6\u914d\u7f6e\u9009\u62e9\u5408\u9002\u7684JAX\u7248\u672c\u3002<\/li>\n<\/ul>\n<\/li>\n<\/ul>\n\n\n\n
            pip <span class<\/strong>=\"token function\">install<\/span> jupyter numpy <span class<\/strong>=\"token string\">\"jax[cuda12]\"<\/span> matplotlib pillow<\/pre>\n\n\n\n
              \n
            • \u4f7f\u7528JAX\u7684\u4e3b\u8981\u529f\u80fd<\/strong>\uff1a\n
                \n
              • \u81ea\u52a8\u5fae\u5206<\/strong>\uff1a\u4f7f\u7528jax.grad<\/code>\u81ea\u52a8\u8ba1\u7b97\u51fd\u6570\u7684\u68af\u5ea6\u3002<\/li>\n<\/ul>\n<\/li>\n<\/ul>\n\n\n\n
                <span class<\/strong>=\"token keyword\">import<\/span> jax<span class<\/strong>=\"token keyword\">import<\/span> jax<span class<\/strong>=\"token punctuation\">.<\/span>numpy <span class<\/strong>=\"token keyword\">as<\/span> jnp<span class<\/strong>=\"token keyword\">def<\/strong><\/span> <span class<\/strong>=\"token function\">cubic_sum<\/span><span class<\/strong>=\"token punctuation\">(<\/span>x<span class<\/strong>=\"token punctuation\">)<\/span><span class<\/strong>=\"token punctuation\">:<\/span>    <span class<\/strong>=\"token keyword\">return<\/strong><\/span> jnp<span class<\/strong>=\"token punctuation\">.<\/span><span class<\/strong>=\"token builtin\">sum<\/span><span class<\/strong>=\"token punctuation\">(<\/span>x<span class<\/strong>=\"token operator\">**<\/span><span class<\/strong>=\"token number\">3<\/span><span class<\/strong>=\"token punctuation\">)<\/span>grad_cubic_sum <span class<\/strong>=\"token operator\">=<\/span> jax<span class<\/strong>=\"token punctuation\">.<\/span>grad<span class<\/strong>=\"token punctuation\">(<\/span>cubic_sum<span class<\/strong>=\"token punctuation\">)<\/span>x_input <span class<\/strong>=\"token operator\">=<\/span> jnp<span class<\/strong>=\"token punctuation\">.<\/span>arange<span class<\/strong>=\"token punctuation\">(<\/span><span class<\/strong>=\"token number\">1.0<\/span><span class<\/strong>=\"token punctuation\">,<\/span> <span class<\/strong>=\"token number\">5.0<\/span><span class<\/strong>=\"token punctuation\">)<\/span>gradient <span class<\/strong>=\"token operator\">=<\/span> grad_cubic_sum<span class<\/strong>=\"token punctuation\">(<\/span>x_input<span class<\/strong>=\"token punctuation\">)<\/span><span class<\/strong>=\"token keyword\">print<\/span><span class<\/strong>=\"token punctuation\">(<\/span><span class<\/strong>=\"token string\">\"\u68af\u5ea6 df\/dx:\"<\/span><span class<\/strong>=\"token punctuation\">,<\/span> gradient<span class<\/strong>=\"token punctuation\">)<\/span><\/pre>\n\n\n\n
                  \n
                • \n
                    \n
                  • \u5373\u65f6\u7f16\u8bd1\uff08JIT\uff09<\/strong>\uff1a\u7528jax.jit<\/code>\u5c06\u51fd\u6570\u7f16\u8bd1\u6210\u4f18\u5316\u540e\u7684\u673a\u5668\u4ee3\u7801\u3002<\/li>\n<\/ul>\n<\/li>\n<\/ul>\n\n\n\n
                    <span class<\/strong>=\"token decorator annotation punctuation\">@jax<span class<\/strong>=\"token punctuation\">.<\/span>jit<\/span><span class<\/strong>=\"token keyword\">def<\/strong><\/span> <span class<\/strong>=\"token function\">selu_jax_jit<\/span><span class<\/strong>=\"token punctuation\">(<\/span>x<span class<\/strong>=\"token punctuation\">)<\/span><span class<\/strong>=\"token punctuation\">:<\/span>    <span class<\/strong>=\"token keyword\">return<\/strong><\/span> <span class<\/strong>=\"token number\">1.0507<\/span> <span class<\/strong>=\"token operator\">*<\/span> jnp<span class<\/strong>=\"token punctuation\">.<\/span>where<span class<\/strong>=\"token punctuation\">(<\/span>x <span class<\/strong>=\"token operator\">><\/span> <span class<\/strong>=\"token number\">0<\/span><span class<\/strong>=\"token punctuation\">,<\/span> x<span class<\/strong>=\"token punctuation\">,<\/span> <span class<\/strong>=\"token number\">1.67326<\/span> <span class<\/strong>=\"token operator\">*<\/span> jnp<span class<\/strong>=\"token punctuation\">.<\/span>exp<span class<\/strong>=\"token punctuation\">(<\/span>x<span class<\/strong>=\"token punctuation\">)<\/span> <span class<\/strong>=\"token operator\">-<\/span> <span class<\/strong>=\"token number\">1.67326<\/span><span class<\/strong>=\"token punctuation\">)<\/span>x_jax <span class<\/strong>=\"token operator\">=<\/span> jnp<span class<\/strong>=\"token punctuation\">.<\/span>random<span class<\/strong>=\"token punctuation\">.<\/span>normal<span class<\/strong>=\"token punctuation\">(<\/span>jax<span class<\/strong>=\"token punctuation\">.<\/span>random<span class<\/strong>=\"token punctuation\">.<\/span>PRNGKey<span class<\/strong>=\"token punctuation\">(<\/span><span class<\/strong>=\"token number\">0<\/span><span class<\/strong>=\"token punctuation\">)<\/span><span class<\/strong>=\"token punctuation\">,<\/span> <span class<\/strong>=\"token punctuation\">(<\/span><span class<\/strong>=\"token number\">10000<\/span><span class<\/strong>=\"token punctuation\">,<\/span> <span class<\/strong>=\"token number\">10000<\/span><span class<\/strong>=\"token punctuation\">)<\/span><span class<\/strong>=\"token punctuation\">)<\/span>result_jax_jit <span class<\/strong>=\"token operator\">=<\/span> selu_jax_jit<span class<\/strong>=\"token punctuation\">(<\/span>x_jax<span class<\/strong>=\"token punctuation\">)<\/span><\/pre>\n\n\n\n
                      \n
                    • \n
                        \n
                      • \u5411\u91cf\u5316<\/strong>\uff1a\u4f7f\u7528jax.vmap<\/code>\u81ea\u52a8\u5c06\u51fd\u6570\u5411\u91cf\u5316\u3002<\/li>\n<\/ul>\n<\/li>\n<\/ul>\n\n\n\n
                        <span class<\/strong>=\"token keyword\">def<\/strong><\/span> <span class<\/strong>=\"token function\">mat_vec_product<\/span><span class<\/strong>=\"token punctuation\">(<\/span>matrix<span class<\/strong>=\"token punctuation\">,<\/span> vector<span class<\/strong>=\"token punctuation\">)<\/span><span class<\/strong>=\"token punctuation\">:<\/span>    <span class<\/strong>=\"token keyword\">return<\/strong><\/span> jnp<span class<\/strong>=\"token punctuation\">.<\/span>dot<span class<\/strong>=\"token punctuation\">(<\/span>matrix<span class<\/strong>=\"token punctuation\">,<\/span> vector<span class<\/strong>=\"token punctuation\">)<\/span>batched_mat_vec <span class<\/strong>=\"token operator\">=<\/span> jax<span class<\/strong>=\"token punctuation\">.<\/span>vmap<span class<\/strong>=\"token punctuation\">(<\/span>mat_vec_product<span class<\/strong>=\"token punctuation\">,<\/span> in_axes<span class<\/strong>=\"token operator\">=<\/span><span class<\/strong>=\"token punctuation\">(<\/span><span class<\/strong>=\"token boolean\">None<\/span><span class<\/strong>=\"token punctuation\">,<\/span> <span class<\/strong>=\"token number\">0<\/span><span class<\/strong>=\"token punctuation\">)<\/span><span class<\/strong>=\"token punctuation\">)<\/span>matrix_jax <span class<\/strong>=\"token operator\">=<\/span> jnp<span class<\/strong>=\"token punctuation\">.<\/span>random<span class<\/strong>=\"token punctuation\">.<\/span>normal<span class<\/strong>=\"token punctuation\">(<\/span>jax<span class<\/strong>=\"token punctuation\">.<\/span>random<span class<\/strong>=\"token punctuation\">.<\/span>PRNGKey<span class<\/strong>=\"token punctuation\">(<\/span><span class<\/strong>=\"token number\">0<\/span><span class<\/strong>=\"token punctuation\">)<\/span><span class<\/strong>=\"token punctuation\">,<\/span> <span class<\/strong>=\"token punctuation\">(<\/span><span class<\/strong>=\"token number\">10000<\/span><span class<\/strong>=\"token punctuation\">,<\/span> <span class<\/strong>=\"token number\">10000<\/span><span class<\/strong>=\"token punctuation\">)<\/span><span class<\/strong>=\"token punctuation\">)<\/span>vectors_jax <span class<\/strong>=\"token operator\">=<\/span> jnp<span class<\/strong>=\"token punctuation\">.<\/span>random<span class<\/strong>=\"token punctuation\">.<\/span>normal<span class<\/strong>=\"token punctuation\">(<\/span>jax<span class<\/strong>=\"token punctuation\">.<\/span>random<span class<\/strong>=\"token punctuation\">.<\/span>PRNGKey<span class<\/strong>=\"token punctuation\">(<\/span><span class<\/strong>=\"token number\">1<\/span><span class<\/strong>=\"token punctuation\">)<\/span><span class<\/strong>=\"token punctuation\">,<\/span> <span class<\/strong>=\"token punctuation\">(<\/span><span class<\/strong>=\"token number\">128<\/span><span class<\/strong>=\"token punctuation\">,<\/span> <span class<\/strong>=\"token number\">10000<\/span><span class<\/strong>=\"token punctuation\">)<\/span><span class<\/strong>=\"token punctuation\">)<\/span>result_vmap <span class<\/strong>=\"token operator\">=<\/span> batched_mat_vec<span class<\/strong>=\"token punctuation\">(<\/span>matrix_jax<span class<\/strong>=\"token punctuation\">,<\/span> vectors_jax<span class<\/strong>=\"token punctuation\">)<\/spa<\/pre>\n\n\n\n
                        JAX\u7684\u5e94\u7528\u573a\u666f<\/h2>\n\n\n\n
                          \n
                        • \u673a\u5668\u5b66\u4e60\u548c\u6df1\u5ea6\u5b66\u4e60<\/strong>\uff1aJAX\u7684\u81ea\u52a8\u5fae\u5206\u548c\u786c\u4ef6\u52a0\u901f\u529f\u80fd\uff0c\u80fd\u9ad8\u6548\u8bad\u7ec3\u548c\u63a8\u7406\u795e\u7ecf\u7f51\u7edc\uff0c\u63d0\u5347\u6a21\u578b\u6027\u80fd\u3002<\/li>\n\n\n\n
                        • \u79d1\u5b66\u8ba1\u7b97<\/strong>\uff1aJAX\u80fd\u8ba1\u7b97\u590d\u6742\u7269\u7406\u65b9\u7a0b\u5bfc\u6570\uff0c\u4f18\u5316\u7269\u7406\u3001\u5316\u5b66\u548c\u6750\u6599\u79d1\u5b66\u4e2d\u7684\u6a21\u62df\u548c\u9884\u6d4b\u3002<\/li>\n\n\n\n
                        • \u6570\u636e\u5206\u6790\u548c\u5904\u7406<\/strong>\uff1a\u501f\u52a9\u5411\u91cf\u5316\u548c\u5e76\u884c\u5316\uff0cJAX\u80fd\u5feb\u901f\u5904\u7406\u5927\u89c4\u6a21\u6570\u636e\uff0c\u9002\u7528\u56fe\u50cf\u3001\u4fe1\u53f7\u5904\u7406\u7b49\u9886\u57df\u3002<\/li>\n\n\n\n
                        • \u91d1\u878d\u5efa\u6a21<\/strong>\uff1a\u7528\u5728\u91d1\u878d\u98ce\u9669\u8bc4\u4f30\u548c\u9ad8\u9891\u4ea4\u6613\uff0c\u9ad8\u6548\u8ba1\u7b97\u52a9\u529b\u5b9e\u65f6\u6570\u636e\u5206\u6790\u548c\u51b3\u7b56\u3002<\/li>\n\n\n\n
                        • \u8ba1\u7b97\u751f\u7269\u5b66<\/strong>\uff1a\u5904\u7406\u57fa\u56e0\u7ec4\u6570\u636e\u3001\u9884\u6d4b\u86cb\u767d\u8d28\u7ed3\u6784\uff0c\u52a0\u901f\u751f\u7269\u533b\u5b66\u7814\u7a76\u548c\u5e94\u7528\u3002<\/li>\n<\/ul>\n","protected":false},"excerpt":{"rendered":"
                          JAX Google\u63a8\u51fa\u7684\u7528\u4e8e\u53d8\u6362\u6570\u503c\u51fd\u6570\u7684\u673a\u5668\u5b66\u4e60\u6846\u67b6 JAX\u662f\u4ec0\u4e48 JAX\u662fGoogle\u63a8\u51fa\u7684\u9ad8\u6027\u80fd\u6570\u503c\u8ba1 […]<\/p>\n","protected":false},"author":1,"featured_media":0,"parent":0,"menu_order":0,"comment_status":"closed","ping_status":"closed","template":"","meta":{"footnotes":""},"class_list":["post-2551","page","type-page","status-publish","hentry"],"_links":{"self":[{"href":"https:\/\/ai.ziyuanzz.online\/index.php\/wp-json\/wp\/v2\/pages\/2551","targetHints":{"allow":["GET"]}}],"collection":[{"href":"https:\/\/ai.ziyuanzz.online\/index.php\/wp-json\/wp\/v2\/pages"}],"about":[{"href":"https:\/\/ai.ziyuanzz.online\/index.php\/wp-json\/wp\/v2\/types\/page"}],"author":[{"embeddable":true,"href":"https:\/\/ai.ziyuanzz.online\/index.php\/wp-json\/wp\/v2\/users\/1"}],"replies":[{"embeddable":true,"href":"https:\/\/ai.ziyuanzz.online\/index.php\/wp-json\/wp\/v2\/comments?post=2551"}],"version-history":[{"count":0,"href":"https:\/\/ai.ziyuanzz.online\/index.php\/wp-json\/wp\/v2\/pages\/2551\/revisions"}],"wp:attachment":[{"href":"https:\/\/ai.ziyuanzz.online\/index.php\/wp-json\/wp\/v2\/media?parent=2551"}],"curies":[{"name":"wp","href":"https:\/\/api.w.org\/{rel}","templated":true}]}}