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

copybara-service merged 3 commits into google:main from samanklesaria:eval_shape_sharding

gemini-code-assist commented on 2026-02-13

IvyZX requested changes on 2026-03-04

Added sharding propagation support in nnx.eval_shape

70351571

Change 'sharding' name to 'out_sharding'

3e3c4594

Support explicit sharding for nnx.eval_shape

6bce1438

samanklesaria force pushed from 4375f506 to 6bce1438 1 day ago

samanklesaria requested a review from

IvyZX 1 day ago

IvyZX approved these changes on 2026-03-06

IvyZX added pull ready

copybara-service merged 49e34d17 into main 21 hours ago

Reviewers

IvyZX

gemini-code-assist

Assignees

No one assigned

Labels

pull ready

Milestone

No milestone