用途
cuequivariance-jax（导入名 cuex）用于在 JAX 中执行 cuequivariance（cue）描述符所定义的等变多项式运算，覆盖 SO3/O3 的 irrep（不可约表示）表示、布局与批处理形状。资料强调应保留项目原有框架与设备：自定义表示使用核心 cuequivariance，PyTorch 模块使用 cuequivariance_torch；本工具仅面向 JAX。
主要能力
- 分段多项式执行：segmented_polynomial 接收描述符多项式、操作数列表与输出 ShapeDtypeStruct 列表。CPU 上使用 method="naive"，支持 JIT、梯度与 vmap。
- 后端与方法：naive 适合便携 JAX 执行（含 CPU）；uniform_1d 与 indexed_linear 属 CUDA 路径，要求每个操作数内分段形状统一、单一模式；indexed_linear 使用 cuex.Repeats，并非普通 CPU 线性层。CPU 路线不需要 GPU 依赖；将数组移到 CPU 不改变方法选择。
- 多项式描述符：如 spherical_harmonics_ir_dict、channelwise_tensor_product、channelwise_tensor_product_ir_dict。channelwise 乘积中权重在前，共享权重无需批轴。filter_irreps_out 只筛选类型，不筛选重数。
- irrep 字典布局：flat_to_dict(layout="ir_mul") 生成 (..., mul, ir.dim) 块；dict_to_flat 不带 layout 参数，输出 mul_ir 布局，需转置各块而非整体 reshape 才能回到 ir_mul。存在重复 irrep 键时须用 RepArray.regroup() 将数据与元数据一起重组。ir_dict.assert_mul_ir_dict 校验块约定，mul_ir_dict 可用模板广播到键。
- Flax NNX 层：cuex.nnx.IrrepsLinear 要求已 regroup 的 irreps 与 nnx.Rngs，按 scale/sqrt(mul_in) 归一化混合同 irrep 通道，未匹配输出为零；IrrepsNormalize(eps, scale, skip_scalars) 按重数平均的平方分量范数归一化；MLP(layer_sizes, activation, dtype, rngs) 使用标量特征与 fan-in 归一化；SphericalHarmonics 在 CPU/GPU 上先归一化向量再调用 naive 多项式，零向量场景需正 eps。
- 密集 RepArray 接口：依据描述符的表示元数据，单输出返回一个 RepArray，.array 持有数据，.reps 映射轴。
使用方式（依据资料示例）
运行前需探测解释器与包：分别尝试导入 jax、cuequivariance、cuequivariance_jax（NNX 任务再加 flax.nnx），打印版本、文件路径与 jax.devices()，复用可用包，只装缺失依赖，并遵守项目 lockfile 或 requirements。资料提示 JAX 0.11.1 / Flax 0.12.10 需要 Python >=3.12；独立组合示例为 jax==0.10.2、cuequivariance==0.13.0、cuequivariance-jax==0.13.0，且该组合也支持 Python 3.11，但不替代项目固定版本。另注明 cuEquivariance 0.13.0 与 JAX 0.12.0 组合执行多项式会报 segmented_polynomial_dce 参数数量错误。
示例脚本先设置 jax_enable_x64，使用 float64，构建 spherical_harmonics_ir_dict 描述符，把向量送入 segmented_polynomial(method="naive")，再与 desc.polynomial 结果用 np.testing.assert_allclose 比对并检查形状。字典示例展示 ir_mul 与 mul_ir 的精确往返转换。密集示例演示广播、jax.jit 与 jax.grad，并用有限差分核对权重梯度。ir_dict 多项式示例按 input_irreps / output_irreps 顺序展开操作数。NNX 示例用 IrrepsLinear 校验输出形状、未匹配 l=2 输出为零及旋转等变性。
资料强调：保存脚本后在独立运行目录执行，轮询同一会话直至退出，读取退出码与数值输出后再报告成功；空输出不能视为断言已运行。安装或修复时不得递归删除目录，隔离环境使用 mktemp 唯一目录。若缺少兼容解释器或可信包源，应报告阻塞与未运行检查。
适用场景与限制
适合编写或调试 cuEquivariance JAX 多项式、irrep 字典、Flax NNX 层，以及在 CPU/CUDA 上验证执行与布局。限制包括：CUDA 路径需兼容的 JAX/GPU/内核构建；indexed-linear 的 CUDA 物种计数须与分组行匹配且总和等于批大小；gather/scatter 需按操作数提供索引（输入在前输出在后，None 表示未索引），重复目标会累加；资料注明不存在 ir_dict.equivariant_polynomial API。具体收费、平台支持范围与更多安装方式暂未提供。