jax.extend.pallas 模块

GridMapping(grid, grid_names, ...[, ...])	GridSpec 的内部规范化版本。
register_lowering_rule(params_cls, rule, ...)