flax
Add sharding propagation support in nnx.eval_shape (clone of #5111)
#5247
Merged

Add sharding propagation support in nnx.eval_shape (clone of #5111) #5247

samanklesaria
gemini-code-assist
gemini-code-assist
gemini-code-assist commented on 2026-02-13
IvyZX
IvyZX requested changes on 2026-03-04
vfdev-5 Added sharding propagation support in nnx.eval_shape
70351571
samanklesaria Change 'sharding' name to 'out_sharding'
3e3c4594
samanklesaria Support explicit sharding for nnx.eval_shape
6bce1438
samanklesaria samanklesaria force pushed from 4375f506 to 6bce1438 1 day ago
samanklesaria samanklesaria requested a review from IvyZX IvyZX 1 day ago
IvyZX
IvyZX approved these changes on 2026-03-06
IvyZX IvyZX added pull ready
copybara-service copybara-service merged 49e34d17 into main 21 hours ago

Login to write a write a comment.

Login via GitHub

Assignees
No one assigned
Labels
Milestone