flax icon indicating copy to clipboard operation
flax copied to clipboard
verified publisher icon google

Make Linen and NNX logical rule deduction align

Open IvyZX opened this issue 4 months ago • 0 comments

  • Moved all the Linen logical axis deduction logic from linen/spmd.py to core/spmd.py, to be shared with NNX APIs.

IvyZX avatar Oct 07 '25 01:10 IvyZX