TensorFlow 如何查询 XLA JIT 编译的算子兼容性【免费下载链接】tensorflowAn Open Source Machine Learning Framework for Everyone项目地址: https://gitcode.com/GitHub_Trending/te/tensorflow如果你的计算图要走 XLA JIT 编译XLA_CPU_JIT 或 XLA_GPU_JIT 编译设备在动手前常需要确认一件事图里用到的算子目标设备是否支持支持的又是哪些数据类型。TensorFlow 仓库为这个场景提供了两条路径直接查tf2xla模块下预生成的算子支持列表或者用tf2xla_supported_ops工具基于当前源码重新生成列表。本文介绍这两条路径的操作方式和输出判断方法。仓库里预生成的两张支持列表列表放在tf2xla的 g3doc 目录下各对应一个 JIT 编译设备cpu_supported_ops.md首行为**Supported operators for device: XLA_CPU_JIT**列出 XLA_CPU_JIT 设备支持的算子gpu_supported_ops.md首行为**Supported operators for device: XLA_GPU_JIT**列出 XLA_GPU_JIT 设备支持的算子。两张表格式相同都是Operator | Type Constraint两列的 Markdown 表格按算子名排序。第二列给出每个类型属性的允许取值例如Conv2D在两表中的约束都是T{float}一行约束为空如NoOp、FFT表示该算子未列出类型约束。如何查一个算子的支持情况按目标设备选表CPU 编译查 cpu_supported_ops.mdGPU 编译查 gpu_supported_ops.md。在表中按算子名查找找到该行即说明该算子被列入了对应设备的 XLA 支持列表表中连名字都没有的算子则未列入该设备的支持列表。对照 Type Constraint 列检查你的实际数据类型。约束按属性名分组例如ArgMax的约束为Tidx{int32,int64}、output_type{int32,int64}、T{float}即只检查与你的输入相关的属性取值是否在列出集合内即可。同一算子在两张表里的约束可能不同判断时不能混用。例如ArgMaxCPU 表为T{float}GPU 表为T{complex64,double,float,int32,int64,uint32,uint64}。再看随机类算子RandomStandardNormal、RandomUniform、RandomUniformInt、TruncatedNormal出现在 CPU 表中而 GPU 表里没有这些行。这类差异正是查表要得到的结论。用 tf2xla_supported_ops 从源码重新生成列表预生成表对应的是生成时的源码。如果你在当前源码树上做了改动、或想确认手头这份代码的实时状态仓库文档在每张表末尾都给出了再生成命令bazel run -c opt -- tensorflow/compiler/tf2xla:tf2xla_supported_ops --deviceXLA_CPU_JITbazel run -c opt -- tensorflow/compiler/tf2xla:tf2xla_supported_ops --deviceXLA_GPU_JIT--device指定要打印哪个编译设备的支持算子取值必须是已注册的编译设备名两张预生成表分别使用XLA_CPU_JIT和XLA_GPU_JIT。该命令只在 TensorFlow 源码根目录下运行目标路径tensorflow/compiler/tf2xla:tf2xla_supported_ops相对于仓库根目录-c opt是文档给出的构建配置。文档未再说明其他环境要求构建所需的环境请按源码树的常规构建流程准备。工具实现见 tf2xla_supported_ops_main.cc 与 tf2xla_supported_ops.cc它注册当前源码的编译内核compilation kernels把算子表打印到标准输出。输出判断与报错行为运行后标准输出的开头格式与仓库两张 md 文档一致首行是**Supported operators for device: XLA_CPU_JIT**设备名随--device变化随后是Operator | Type Constraint表头和按算子名排序的表格结尾还会再打印一次再生成命令。把这份输出与仓库中的对应文档交叉对照即可确认当前源码与支持表是否一致。两个报错场景由源码中的检查逻辑给出传入未注册的设备名时工具报Unknown device: 名称并打印用法该命令只接受 flags不接受其他位置参数多余参数会触发错误提示。使用边界两张表描述的是 XLA_CPU_JIT / XLA_GPU_JIT 这两个 JIT 编译设备的支持列表判断结果只对相应设备有效不要拿 CPU 表的结论用于 GPU 或反之。表中列出的是该设备支持的算子参考文档未说明未列入列表的算子在运行时如何处理本文也不做推断。预生成表随源码版本更新查表前先确认文档对应你正在使用的源码版本有本地改动时以bazel run的输出为准。【免费下载链接】tensorflowAn Open Source Machine Learning Framework for Everyone项目地址: https://gitcode.com/GitHub_Trending/te/tensorflow创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考