Fix custom getter handling in tpu.rewrite() and friends.
It used to save the existing custom getter then overwrites the custom getter. That means the previous custom getter will never be called inside "computation". It now create a new custom getter that calls the previous custom getter. PiperOrigin-RevId: 214715720
Loading
Please sign in to comment