capture_tpu_profile
PiperOrigin-RevId: 340943636 Change-Id: I8347dcdf63d38ee0368ad521082441555e2b3222