JAX 原语数据类型覆盖率指南解读 jax2tf 的 133 个数值原语与 7554 个测试装置【免费下载链接】jaxComposable transformations of PythonNumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more项目地址: https://gitcode.com/GitHub_Trending/ja/jax本指南以 JAX 仓库中的自动生成文档 jax_primitives_coverage.md 为骨架系统讲解 JAX 数值原语primitive在各类数据类型dtype与设备CPU/GPU/TPU上的支持边界包括如何读懂覆盖率表中的 dtype 缩写、133 个原语的完整覆盖情况、13 类部分实现限制及其成因以及如何借助 jax_primitives_coverage_test.py 与 test_harnesses.py 从源码层面重新生成并维护这份覆盖率清单。读完本文你将能够快速判断某个 JAX 原语在目标 dtype/设备组合下是否可用并掌握自行验证、更新覆盖率文档的完整方法。这份文档是什么自动生成的原语能力边界jax_primitives_coverage.md是一份自动生成的清单用于记录 jax2tf 与 JAX 内部测试所关心的数值原语支持情况。文档头部标注了生成日期Last generated on: 2023-07-31并声明如果出现新的限制自动化测试会失败但限制被修复后测试不会自动失败。因此若读者发现某个已记录的限制不再存在应提请更新本文件。其生成机制非常明确文档由 jax_primitives_coverage.md.template 模板填充而来模板中预留了{{generation_date}}、{{nr_harnesses}}、{{nr_primitives}}、{{primitive_coverage_table}}、{{primitive_unimpl_table}}五个占位符。当前版本的文档基于7554 个测试装置test harnesses验证133 个数值 JAX 原语的实现。测试装置harness是一个关键概念它定义在 test_harnesses.py 中每个Harness封装了一个可调用对象通常是lax下的某个原语、一组参数描述符RandArg随机数组、StaticArg静态参数等、一个目标 dtype以及一组jax_unimplemented限制Limitation。文档中那句这实际上展示的是每个原语在任何设备上都没有被当前测试装置覆盖的 dtype准确点出了第一张表的本质——它是测试覆盖范围的记录而非绝对的数学性质声明。读懂 dtype 缩写表前的八组集合记号为了压缩表格宽度文档定义了一组 dtype 集合缩写它们被用于两张覆盖表的支持/不支持列中。展开如下缩写含义包含的 dtypesignedint8、int16、int32、int64unsigneduint8、uint16、uint32、uint64integersignedunsignedfloatingfloat16、bfloat16、float32、float64complexcomplex64、complex128inexactfloatingcomplexallintegerinexactbool这套缩写的实际展开逻辑可以在源码 test_harnesses.py 的dtypes_to_str函数中找到它对一组 dtype 名称做集合运算当完整包含某个缩写对应的集合时就用缩写替代例如同时包含int8/int16/int32/int64就折叠为signed最终按字典序输出。这正是全覆盖类单元格显示为all、而缺个别类型时显示为bool, unsigned这类组合的原因。另一个必须注意的前提64 位类型float64、int64、uint64、complex128只有设置了--jax_enable_x64标志或JAX_ENABLE_X64环境变量时才被纳入支持的 dtype集合。也就是说默认的 32 位模式下文档中标记为signed/floating/complex的覆盖并不包含对应的 64 位变体。表一133 个原语在各 dtype 上的支持全貌判定规则一个原语只要在至少一种设备类型上受支持就认为该 dtype 被支持。下表完整列出每个原语的总测试装置数、至少在一种设备上受支持的 dtype以及在任何设备上都未被测试的 dtype。PrimitiveTotal test harnessesdtypes supported on at least one devicedtypes NOT tested on any deviceabs10inexact, signedbool, unsignedacos6inexactbool, integeracosh6inexactbool, integeradd16inexact, integerbooladd_any14inexact, integerbooland11bool, integerinexactapprox_top_k24floatingbool, complex, integerargmax64bool, floating, integercomplexargmin64bool, floating, integercomplexasin6inexactbool, integerasinh6inexactbool, integeratan6inexactbool, integeratan26floatingbool, complex, integeratanh6inexactbool, integerbessel_i0e4floatingbool, complex, integerbessel_i1e4floatingbool, complex, integerbitcast_convert_type41allbroadcast_in_dim19allcbrt4floatingbool, complex, integerceil4floatingbool, complex, integercholesky30inexactbool, integerclamp20allcomplex4float32, float64bfloat16, bool, complex, float16, integerconcatenate17allconj5complex, float32, float64bfloat16, bool, float16, integerconv_general_dilated132inexact, signedbool, unsignedconvert_element_type201allcos6inexactbool, integercosh6inexactbool, integercumlogsumexp12float16, float32, float64bfloat16, bool, complex, integercummax34inexact, integerboolcummin34inexact, integerboolcumprod34inexact, integerboolcumsum34inexact, integerboolcustom_linear_solve4float32, float64bfloat16, bool, complex, float16, integerdevice_put16alldigamma4floatingbool, complex, integerdiv20inexact, integerbooldot_general400alldynamic_slice68alldynamic_update_slice46alleig72inexactbool, integereigh36inexactbool, integereq17allerf4floatingbool, complex, integererf_inv4floatingbool, complex, integererfc4floatingbool, complex, integerexp6inexactbool, integerexpm16inexactbool, integerfft32complex, float32, float64bfloat16, bool, float16, integerfloor4floatingbool, complex, integergather164allge17allgt17alligamma6floatingbool, complex, integerigammac6floatingbool, complex, integerimag2complexbool, floating, integerinteger_pow108inexact, integerbooliota16inexact, integerbooliota_2x32_shape3uint32bool, inexact, signed, uint16, uint64, uint8is_finite4floatingbool, complex, integerle17alllgamma4floatingbool, complex, integerlog6inexactbool, integerlog1p6inexactbool, integerlogistic6inexactbool, integerlt17alllu18inexactbool, integermax27allmin27allmul16inexact, integerboolne17allneg14inexact, integerboolnextafter6floatingbool, complex, integeror11bool, integerinexactpad180allpopulation_count8integerbool, inexactpow10inexactbool, integerqr60inexactbool, integerrandom_categorical12floatingbool, complex, integerrandom_gamma4float32, float64bfloat16, bool, complex, float16, integerrandom_randint12signedbool, inexact, unsignedrandom_split5uint32allrandom_uniform12floatingbool, complex, integerreal2complexbool, floating, integerreduce33allreduce_and1boolinexact, integerreduce_max15allreduce_min15allreduce_or1boolinexact, integerreduce_precision32floatingbool, complex, integerreduce_prod14inexact, integerboolreduce_sum14inexact, integerboolreduce_window_add50inexact, integerboolreduce_window_max66allreduce_window_min27allreduce_window_mul42inexact, integerboolregularized_incomplete_beta4floatingbool, complex, integerrem18floating, integerbool, complexreshape19allrev19allrng_bit_generator36uint32, uint64bool, inexact, signed, uint16, uint8round6floatingbool, complex, integerrsqrt6inexactbool, integerscatter645allscatter_add885allscatter_max885allscatter_min888allscatter_mul885allselect_and_gather_add15floatingbool, complex, integerselect_and_scatter_add27bool, floating, integercomplexselect_n32allshift_left10integerbool, inexactshift_right_arithmetic10integerbool, inexactshift_right_logical10integerbool, inexactsign28inexact, integerboolsin6inexactbool, integersinh6inexactbool, integerslice24allsort29allsqrt6inexactbool, integersqueeze23allstop_gradient15allsub16inexact, integerboolsvd120inexactbool, integertan6inexactbool, integertanh6inexactbool, integertop_k15bool, floating, integercomplextranspose17alltriangular_solve26inexactbool, integertridiagonal_solve2float32, float64bfloat16, bool, complex, float16, integerxor11bool, integerinexactzeros_like15all从表一能得出的几类结论全覆盖原语约 40 个原语的未测试列为空即所有 dtypeall都在至少一种设备上被测试过例如bitcast_convert_type、broadcast_in_dim、convert_element_type、dot_general400 个测试装置表中最多、gather、scatter系列、pad、reduce系列、sort等。这类原语通常是布局/搬运/规约类操作对 dtype 几乎无限制。布尔与整数的特殊位置数值函数如sin、exp、log天然不支持bool/integer而位运算and、or、xor只支持bool, integer不支持inexact。population_count、shift_left/right_*则只支持integer。窄覆盖原语iota_2x32_shape仅在uint32上受测random_split仅在uint32上受测rng_bit_generator覆盖uint32, uint64——它们服务于随机数/键生成内部机制覆盖面刻意很窄。complex、conj、real、imag等复数相关原语仅覆盖复数与 64 位浮点fft覆盖complex, float32, float64均不接受float16/bfloat16。表二13 类部分实现的 dtype/设备限制即使某个 dtype 在某设备上被支持也可能只是部分实现要么只有部分设备实现了该路径要么存在 XLA 编译器尚未实现的组合。表二记录的就是这类限制该表已剔除表一中完全不支持的 dtypeAffected primitiveDescription of limitationAffected dtypesAffected devicescholeskyunimplementedfloat16cpu, gpuclampunimplementedbool, complexcpu, gpu, tpuconv_general_dilatedpreferred_element_type not implemented for integerssignedgpudot_generalpreferred_element_type must be floating for integer dtypeintegergpudot_generalpreferred_element_type must match dtype for floating pointinexactgpueigonly supported on CPU in JAXalltpu, gpueigunimplementedbfloat16, float16cpueighunimplementedbfloat16, float16cpu, gpuluunimplementedbfloat16, float16cpu, gpu, tpuqrunimplementedbfloat16, float16cpu, gpuscatter_addunimplementedboolcpu, gpu, tpuscatter_mulunimplementedboolcpu, gpu, tpuselect_and_scatter_addworks only for 2 or more inactive dimensionsalltpusvdunimplementedbfloat16, float16cpu, gputriangular_solveunimplementedfloat16gpu文档特别给出了两个典型的部分实现成因示例lax.eig特征分解JAX 只在CPU 和 GPU上用自定义内核实现了它没有 TPU 实现——这是设备实现的真实缺口。lu、qr、eigh、svd等线性代数原语对bfloat16、float16普遍unimplemented因为底层 LAPACK/定制内核在这些低精度类型上没有对应入口。dot_general/conv_general_dilated在 GPU 上对preferred_element_type参数有额外约束——整数 dtype 要求preferred_element_type必须是浮点浮点 dtype 则要求它与输入 dtype 一致。scatter_add/scatter_mul在全部三种设备上都不支持bool归约这个限制在 test_harnesses.py 中通过Limitation(unimplemented, dtypes[np.bool_], enabled(f_lax in [lax.scatter_add, lax.scatter_mul]))显式声明——enabled参数精确控制了该限制只作用于scatter_add与scatter_mul两个原语。底层机制Harness、Limitation 与表格生成逻辑测试装置Harness与限制Limitation的声明式结构test_harnesses.py 中的Harness类是整套覆盖率体系的载体。其核心字段包括group_name分组名通常是原语名经sanitize_test_name清洗name装置描述名与group_name拼接成fullname用于测试用例命名fun被测试的可调用对象如lax.slicearg_descriptors参数描述符序列支持RandArg(shape, dtype)随机张量、StaticArg(value)静态参数如轴号、CustomArg(fun)dtype该装置覆盖的 dtypejax_unimplemented一组Limitation描述该装置在 JAX 中未实现/部分实现的情形。Limitationtest_harnesses.py用description、enabled、devices、dtypes、skip_run五个字段编码限制条件其filter(device, dtype)方法负责判断某限制是否命中当前设备 当前 dtype组合。运行测试时test_jax_implemented若执行装置抛出的异常恰好命中一个已声明且被过滤命中的限制测试视为通过反之若未命中任何限制却抛错测试失败。这也解释了文档的声明新的限制会使测试失败但修复后的限制不会自动使测试失败——因为代码中找到太多限制的断言是被注释掉的jax_primitives_coverage_test.py而找到太少限制即意外报错会直接失败。覆盖率文档的生成管线jax_primitives_coverage_test.py 中的test_generate_primitives_coverage_doc完整演示了两张表的生成过程遍历test_harnesses.all_harnesses按group_name分组统计每个原语的装置数量并合并每个装置测试过的 dtype至少在一种设备上受测的集合用all_dtypes - dtypes_tested计算未在任何设备上测试的 dtype 集合交给dtypes_to_str缩写格式化后写入第一张表收集所有enabled的Limitation按(group_name, description, devices, dtypes)去重后写入第二张表读取模板 jax_primitives_coverage.md.template替换五个占位符后写回jax_primitives_coverage.md。注意第二步的细节文档中dtypes supported on at least one device列实际是所有装置 dtype 的并集因此装置数量越多的原语如scatter_min的 888 个、scatter_add的 885 个其覆盖率越全面。重新生成覆盖率表命令与前置条件文档给出了在CPU 机器上重新生成两张表的完整命令这也是测试代码中强制的环境约束JAX_OUTPUT_LIMITATIONS_DOC1 JAX_ENABLE_X641 python jax/experimental/jax2tf/tests/jax_primitives_coverage_test.py JaxPrimitiveTest.test_generate_primitives_coverage_doc两个环境变量的作用如下环境变量作用JAX_OUTPUT_LIMITATIONS_DOC1触发文档输出。未设置时测试直接跳过jax_primitives_coverage_test.py 中raise unittest.SkipTestJAX_ENABLE_X641开启 64 位类型否则 64 位 dtype 无法进入支持集合表格会失真此外测试代码还断言当前设备必须是cpu文档注释说明 CPU/GPU 支持的 dtype 比 TPU 更多因此只能在 CPU 上生成以保证完整性且enable_x64配置必须为真jax_primitives_coverage_test.py。若只想调试单个装置的测试用例可在 jax_primitives_coverage_test.py 的test_harnesses.parameterized(...)装饰器中启用one_containingfoo参数或设置JAX_TEST_HARNESS_ONE_CONTAINING环境变量test_harnesses.py。排查实际限制如何定位具体报错当你在自己的设备上遇到某个原语在某 dtype 下不可用想确认是否为已知限制时运行test_jax_implemented测试它会对全部装置执行 JAX 调用并核对jax_unimplemented字段在测试日志中搜索 limitation关键字即可看到所有命中与未命中的限制描述、涉及装置与 dtypejax_primitives_coverage_test.py对照本文两张表先查表一确认该 dtype 是否至少在一种设备上受支持再查表二确认目标设备是否在部分实现的受影响列表中若发现表中记录的限制已不存在按文档建议提请更新 jax_primitives_coverage.md——因为这类已修复的限制不会被测试自动发现。与 jax2tf 限制体系的关联本覆盖率文档描述的是JAX 原生实现的限制jax_unimplemented而 jax2tf 转换到 TensorFlow 时还存在另一层独立的限制体系——jax2tf_limitations.py 中的Jax2TfLimitation继承自同一个Limitation基类但额外增加了modeseager、graph、compiled、expect_tf_error、custom_assert、tol容差等 jax2tf 特有的字段用于刻画转换后代码在 TF 三种执行模式下的数值偏差与失败情形。两者共同构成了JAX 原语 → 各设备 / TensorFlow两条路径上的完整兼容性地图前者回答这个原语在 JAX 里能用吗后者回答转换到 TF 后还能用吗、误差多大。小结与使用建议判断原语 × dtype 是否可用时以表一为第一层过滤未被任何设备测试的 dtype 组合默认视为不可用再以表二为第二层过滤确认目标设备不在部分实现限制列表中。表一的支持语义是至少一种设备跨设备移植代码前务必再核对表二中的设备列。涉及 64 位类型时务必确认运行环境已开启JAX_ENABLE_X64否则表一中的signed/floating/complex覆盖范围会与实际不符。若需要最新数据按上文命令在 CPU 环境重新生成文档若只是排查运行时报错从test_jax_implemented的日志中搜索 limitation 即可获得与当前设备最相关的精确信息。【免费下载链接】jaxComposable transformations of PythonNumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more项目地址: https://gitcode.com/GitHub_Trending/ja/jax创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考