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) 132 days ago

vfdev-5 commented on 2025-09-16

samanklesaria marked this pull request as ready for review 130 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` 130 days ago

samanklesaria marked this pull request as draft 130 days ago

samanklesaria force pushed from 172d1fd5 to c02e9f37 130 days ago

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

samanklesaria requested a review from

vfdev-5 130 days ago

vfdev-5 commented on 2025-09-18

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

samanklesaria requested a review from

vfdev-5 129 days ago

vfdev-5 commented on 2025-09-19

samanklesaria force pushed from 44d88609 to c9b10fcb 128 days ago

vfdev-5 commented on 2025-09-22

vfdev-5 commented on 2025-09-23

samanklesaria requested a review from

vfdev-5 124 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 124 days ago

cgarciae commented on 2025-09-23

Add check for duplicate nodes

44ed2fd9

samanklesaria requested a review from

vfdev-5 123 days ago

Support shared structure

126b2139

samanklesaria force pushed from 93b1066a to 126b2139 122 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 118 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