try to make lstm composite op

Author: jakesabathia2

coreai_torch/_aten_to_core.py

@@ -56,6 +56,7 @@
 from ._utils import (
     replace_pad_with_mode as _replace_pad_with_mode,
 )
+from .composite_ops._lstm import replace_lstm
 
 INT32_MAX: int = 2147483647
 
@@ -3640,6 +3641,7 @@ def sdpa_maskless(q: Value, k: Value, v: Value) -> Value:
     "log.default": replace_unary_ops,
     "log10.default": replace_log10,
     "log1p.default": replace_log1p,
+    "lstm.input": replace_lstm,
     "log2.default": replace_log2,
     "logical_and.default": replace_binary_logical_ops,
     "logical_not.default": replace_logical_not,

coreai_torch/_decomp.py

@@ -17,6 +17,7 @@
     torch.ops.aten.hardsigmoid.default,
     torch.ops.aten.hardswish.default,
     torch.ops.aten.instance_norm.default,
+    torch.ops.aten.lstm.input,
     torch.ops.aten.pixel_shuffle.default,
     torch.ops.aten.reflection_pad1d.default,
     torch.ops.aten.reflection_pad2d.default,
@@ -40,6 +41,7 @@ def get_decomp_table() -> dict:
 
     * ``torch.ops.aten.hardsigmoid.default``
     * ``torch.ops.aten.instance_norm.default``
+    * ``torch.ops.aten.lstm.input`` (lowered to a ``"lstm"`` composite op)
     * ``torch.ops.aten.pixel_shuffle.default``
     * ``torch.ops.aten.scaled_dot_product_attention.default``