Update the flops computation by counting add and multiply separately.

PiperOrigin-RevId: 534193569

Author: The paxml Authors

paxml/tools/model_analysis.py

@@ -30,7 +30,7 @@
 
 ##### bert.BertAdamL4H128 #####
 
-GFLOPS = 23.79
+GFLOPS = 47.58
 
 
 **************
@@ -174,7 +174,7 @@ def get_flops(self, fprop_func: str = 'compute_predictions'):
       client = jax.lib.xla_bridge.get_backend()
       m = jax.xla_computation(model_fprop)(datum).as_hlo_module()
       analysis = jax.lib.xla_client._xla.hlo_module_cost_analysis(client, m)
-      flops = analysis['flops'] / 2.0
+      flops = analysis['flops']
       gflops = flops / 1e9
 
       print('\n' + '#' * 5 + ' ' + self.exp_name + ' ' + '#' * 5)