用途：cuequivariance_jax（导入名 cuex）在 GPU 上通过 JAX 执行 cuequivariance 等变多项式，面向等变图神经网络等场景。
主要能力：核心原语 cuex.segmented_polynomial() 是支持完整 AD、vmap、JIT 的 JAX 原语。提供两种数据表示：cuex.equivariant_polynomial() 配合 RepArray（带表示元数据的连续数组），以及 cuex.ir_dict 模块的 dict[Irrep, Array] 接口（使用 IrDictPolynomial 描述符，可与 jax.tree 自然配合）。另有 cuex.nnx 模块提供 Flax NNX 层封装。
执行方法："naive" 为纯 JAX，任意平台可用；"uniform_1d" 为 CUDA 内核，需 GPU 且每个操作数内各段形状一致、单一模式；"indexed_linear" 为 CUDA 内核，用于带 cuex.Repeats 索引的线性运算。
使用方式：先由 cue.descriptors 构建描述符并取 .polynomial，再调用 cuex.segmented_polynomial(poly, inputs, output_specs, method=...)。输入可含任意数量批轴并遵循 NumPy 广播；索引数组提供 gather（输入）与 scatter（输出），每个操作数一个索引，None 表示不索引。ir_dict 接口使用 _ir_dict 描述符变体，配合 segmented_polynomial_uniform_1d 及 input_indices/output_indices 完成消息传递。NNX 层包括 IrrepsLinear、SphericalHarmonics、IrrepsNormalize、MLP（仅标量）与 IrrepsIndexedLinear。
适用场景：NequIP、MACE 类等变 GNN 的张量积与消息传递；按 irrep 拆分操作数后，可与逐 irrep 线性层直接衔接，无需转置或拷贝。
限制：uniform_1d 与 indexed_linear 需 GPU；uniform_1d 要求各段形状一致且单一模式；IrrepsLinear 的 irreps 需 regroup；MLP 仅支持标量。安装命令、版本与收费信息暂未提供。