flax
Add compute_flops and compute_vjp_flops options to `nnx.tabulate`
#4948

Merged

Add compute_flops and compute_vjp_flops options to `nnx.tabulate` #4948

copybara-service merged 4 commits into google:main from samanklesaria:issues/4883

samanklesaria changed the title ~~Add compute_flops and compute_vjp_flops options to `nnx.tabulate`.~~ Add compute_flops and compute_vjp_flops options to `nnx.tabulate` (WIP) 106 days ago

vfdev-5 commented on 2025-09-16

samanklesaria marked this pull request as ready for review 104 days ago

samanklesaria changed the title ~~Add compute_flops and compute_vjp_flops options to `nnx.tabulate` (WIP)~~ Add compute_flops and compute_vjp_flops options to `nnx.tabulate` 104 days ago

samanklesaria marked this pull request as draft 104 days ago

samanklesaria force pushed from 172d1fd5 to c02e9f37 104 days ago

samanklesaria marked this pull request as ready for review 104 days ago

samanklesaria requested a review from

vfdev-5 104 days ago

vfdev-5 commented on 2025-09-18

samanklesaria force pushed from 39b15016 to 44d88609 104 days ago

samanklesaria requested a review from

vfdev-5 103 days ago

vfdev-5 commented on 2025-09-19

samanklesaria force pushed from 44d88609 to c9b10fcb 103 days ago

vfdev-5 commented on 2025-09-22

vfdev-5 commented on 2025-09-23

samanklesaria requested a review from

vfdev-5 99 days ago

vfdev-5 commented on 2025-09-23

vfdev-5 approved these changes on 2025-09-23

Add compute_flops and compute_vjp options to tabulate

27ed2531

samanklesaria force pushed from 0efffe90 to 27ed2531 98 days ago

cgarciae commented on 2025-09-23

Add check for duplicate nodes

44ed2fd9

samanklesaria requested a review from

vfdev-5 98 days ago

Support shared structure

126b2139

samanklesaria force pushed from 93b1066a to 126b2139 96 days ago

vfdev-5 approved these changes on 2025-09-25

vfdev-5 added pull ready

Remove type annotation

8b7c8bd5

copybara-service merged ebb980a6 into main 92 days ago

Reviewers

vfdev-5

cgarciae

Assignees

No one assigned

Labels

pull ready

Milestone

No milestone

flax Add compute_flops and compute_vjp_flops options to `nnx.tabulate` #4948 Merged

Add compute_flops and compute_vjp_flops options to `nnx.tabulate` #4948

flax
Add compute_flops and compute_vjp_flops options to `nnx.tabulate`
#4948

Merged