flax.linen.with\_partitioning
=============================

.. currentmodule:: flax.linen

.. autofunction:: with_partitioning