[llvm] 3486b1b - [mlgo][nfc] regalloc test model generator: prep for TFLite

[Mircea Trofin via llvm-commits]

Author: Mircea Trofin
Date: 2022-08-11T15:53:23-07:00
New Revision: 3486b1b7362eb30c70aa913a06933dad632f7964

URL: https://github.com/llvm/llvm-project/commit/3486b1b7362eb30c70aa913a06933dad632f7964
DIFF: https://github.com/llvm/llvm-project/commit/3486b1b7362eb30c70aa913a06933dad632f7964.diff

LOG: [mlgo][nfc] regalloc test model generator: prep for TFLite

Casting operator to make TFLite happy.

Reviewed By: yundiqian

Differential Revision: https://reviews.llvm.org/D131584

Added: 
    

Modified: 
    llvm/lib/Analysis/models/gen-regalloc-eviction-test-model.py

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Analysis/models/gen-regalloc-eviction-test-model.py b/llvm/lib/Analysis/models/gen-regalloc-eviction-test-model.py
index 11bc3f259ddee..e41e71a09d828 100644
--- a/llvm/lib/Analysis/models/gen-regalloc-eviction-test-model.py
+++ b/llvm/lib/Analysis/models/gen-regalloc-eviction-test-model.py
@@ -46,7 +46,8 @@ def build_mock_model(path):
   module.var = tf.Variable(0, dtype=tf.int64)
 
   def action(*inputs):
-    result = tf.math.argmax(tf.cast(inputs[0]['mask'], tf.int32), axis=-1) + module.var
+    result = tf.math.argmax(
+        tf.cast(inputs[0]['mask'], tf.int32), axis=-1) + module.var
     return {POLICY_DECISION_LABEL: result}
   module.action = tf.function()(action)
   action = {