cuequivariance-jax

RepArray

RepArray(reps, array[, layout])

一个用 cue.Rep 字典装饰的 jax.Array，用于表示在群表示下变换的轴。

from_segments(irreps, segments, shape[, ...])	从分段构造一个 RepArray。
as_irreps_array(input[, layout, like])	将输入转换为 RepArray。
concatenate(arrays)	连接一个 cuex.RepArray 列表
randn(key, rep[, leading_shape, dtype])	生成一个随机的 cuex.RepArray。

张量积

equivariant_tensor_product(e, *inputs[, ...])	计算输入数组的等变张量积。
tensor_product(d, *inputs[, dtype_output, ...])	计算 SegmentedTensorProduct 的最后一个操作数。

额外模块

flax_linen.Linear(irreps_out, layout, force, ...)	等变线性层。
flax_linen.LayerNorm(epsilon, parent, ...)

spherical_harmonics(ls, vector[, normalize, ...])

计算一个向量的球谐函数。