摘要：
这篇文章中，我简要实现一下大语言模型的 MOE 模块。MOE 模块位于每个GPT层中，位于注意力模块的后面，每个MOE模块包含若干个MLP模块作为专家。这些专家是稀疏的，也就是每次选择部分来调用，并不会调用全部，从而节省宝贵的算力。 首先定义一些常量，通常应该在模型配置文件里面。 bs = 5 # 阅读全文
摘要：
JAX 中文文档（一） 原文：jax.readthedocs.io/en/latest/ 开始入门 安装 JAX 原文：jax.readthedocs.io/en/latest/installation.html 使用 JAX 需要安装两个包：jax 是纯 Python 的跨平台库，jaxlib 包 阅读全文
摘要：
JAX 中文文档（五） 原文：jax.readthedocs.io/en/latest/ 形状多态性 原文：jax.readthedocs.io/en/latest/export/shape_poly.html 当使用 JIT 模式的 JAX 时，函数将被跟踪、降级到 StableHLO，并针对每种 阅读全文
摘要：
JAX 中文文档（四） 原文：jax.readthedocs.io/en/latest/ 理解 Jaxpr 原文：jax.readthedocs.io/en/latest/jaxpr.html 更新日期：2020 年 5 月 3 日（提交标识为 f1a46fe）。 从概念上讲，可以将 JAX 转换看 阅读全文
摘要：
JAX 中文文档（十一） 原文：jax.readthedocs.io/en/latest/ JAX 增强提案（JEPs） 原文：jax.readthedocs.io/en/latest/jep/index.html 大多数改动可以通过简单的问题/讨论和拉取请求进行讨论。 然而，有些变更范围较大或需要 阅读全文
摘要：
JAX 中文文档（十五） 原文：jax.readthedocs.io/en/latest/ jax.tree 模块 原文：jax.readthedocs.io/en/latest/jax.tree.html 用于处理树形容器数据结构的实用工具。 jax.tree 命名空间包含了来自 jax.tree 阅读全文
摘要：
JAX 中文文档（十四） 原文：jax.readthedocs.io/en/latest/ jax.scipy 模块 原文：jax.readthedocs.io/en/latest/jax.scipy.html jax.scipy.cluster | vq(obs, code_book[, chec 阅读全文
摘要：
JAX 中文文档（十三） 原文：jax.readthedocs.io/en/latest/ 在 JAX 之上构建 原文：jax.readthedocs.io/en/latest/building_on_jax.html 学习高级 JAX 使用的一种很好的方法是看看其他库如何使用 JAX，它们如何将库 阅读全文
摘要：
JAX 中文文档（十七） 原文：jax.readthedocs.io/en/latest/ JAX 术语表 原文：jax.readthedocs.io/en/latest/glossary.html 数组 JAX 的 numpy.ndarray 的类比。见 jax.Array。 CPU 缩写Cent 阅读全文
摘要：
JAX 中文文档（十六） 原文：jax.readthedocs.io/en/latest/ jax.experimental.sparse.bcoo_multiply_dense 原文：jax.readthedocs.io/en/latest/_autosummary/jax.experimenta 阅读全文
摘要：
JAX 中文文档（十二） 原文：jax.readthedocs.io/en/latest/ Jax 和 Jaxlib 版本控制 原文：jax.readthedocs.io/en/latest/jep/9419-jax-versioning.html 为什么 jax 和 jaxlib 是独立的包？ 我 阅读全文
摘要：
JAX 中文文档（十） 原文：jax.readthedocs.io/en/latest/ JAX 中的广义卷积 原文：jax.readthedocs.io/en/latest/notebooks/convolutions.html JAX 提供了多种接口来跨数据计算卷积，包括： jax.numpy. 阅读全文
摘要：
JAX 中文文档（三） 原文：jax.readthedocs.io/en/latest/ 有状态计算 原文：jax.readthedocs.io/en/latest/stateful-computations.html JAX 的转换（如jit()、vmap()、grad()）要求它们包装的函数是纯 阅读全文
摘要：
JAX 中文文档（七） 原文：jax.readthedocs.io/en/latest/ 使用 shard_map 的 SPMD 多设备并行性 原文：jax.readthedocs.io/en/latest/notebooks/shard_map.html shard_map 是一种单程序多数据（S 阅读全文
摘要：
JAX 中文文档（六） 原文：jax.readthedocs.io/en/latest/ 高级教程 原文：jax.readthedocs.io/en/latest/advanced_guide.html 本节包含更高级主题的示例和教程，如多核计算、自定义操作及更深入的应用 示例 使用 tensorf 阅读全文
摘要：
JAX 中文文档（九） 原文：jax.readthedocs.io/en/latest/ 使用jax.checkpoint控制自动微分的保存数值（又名jax.remat） 原文：jax.readthedocs.io/en/latest/notebooks/autodiff_remat.html im 阅读全文
摘要：
JAX 中文文档（二） 原文：jax.readthedocs.io/en/latest/ JAX 教程 原文：jax.readthedocs.io/en/latest/tutorials.html 快速入门 关键概念 即时编译 自动向量化 自动微分 调试入门 伪随机数 使用 pytrees 工作 分 阅读全文
摘要：
JAX 中文文档（八） 原文：jax.readthedocs.io/en/latest/ 自动微分手册 原文：jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html alexbw@, mattjj@ JAX 拥有非常通用的自动微分系 阅读全文