jax

438ed6d6 - Mark `PartitionSpec` as `is_final()` so that it can't be subclassed.