[HARD] 支持任意张量并行分割

上下文: 张量并行允许通过在多个设备之间分发权重将模型分割到多个设备。张量并行的自动并行策略使用 MLX LM 的 shard_linear 和 shard_inplace 函数来均匀地分割张量。这些函数假设分片维度能够整除组大小。