spmd# flax.nnx.get_partition_spec(tree)[源代码]# 从包含 Variable 值的 PyTree 中提取 PartitionSpec 树。 flax.nnx.get_named_sharding(tree, mesh)[源代码]# flax.nnx.with_partitioning(initializer, sharding, mesh=None, **metadata)[源代码]# flax.nnx.with_sharding_constraint(x, axis_resources, mesh=None)[源代码]#
