flax.linen.with_logical_partitioning#

flax.linen.with_logical_partitioning(fn, names)[source]#

Wraps a function’s return value with LogicallyPartitioned.

Example:

kernel_init = with_logical_partitioning(
    nn.initializers.lecun_normal, (None, "data"))
partitioned_dense = nn.Dense(features, kernel_init=kernel_init)
Parameters
  • fn – The function to be wrapped. Typically this is an initializer.

  • names – The logical axis passed to LogicallyPartitioned.

Returns

A function wrapping fn that will return an instance of LogicallyPartitioned.