add the meta registration for choose_tpp_linear_weight (#3020)

intel · Jul 1, 2024 · c8eb8aa · c8eb8aa
1 parent 1d02cee
commit c8eb8aa
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/intel_extension_for_pytorch/_meta_registrations.py b/intel_extension_for_pytorch/_meta_registrations.py
@@ -521,6 +521,16 @@ def meta_tpp_linear_bias(
     return input.new_empty((*input.shape[:-1], out_features))
 
 
+@register_meta("choose_tpp_linear_weight")
+def meta_choose_tpp_linear_weight(x, weight, weight_for_large_batch):
+    M = x.numel() // x.size(-1)
+    return (
+        weight_for_large_batch
+        if weight_for_large_batch is not None and M >= 256
+        else weight
+    )
+
+
 @register_meta("tpp_linear_gelu")
 def meta_tpp_linear_gelu(
     input,