Add impl for aten.view.default, tfl.reshape and lowering.

PiperOrigin-RevId: 743746772

Author: junjiang-lab

ai_edge_torch/odml_torch/experimental/torch_tfl/_decomps.py

@@ -116,6 +116,11 @@ def _aten_permute_decomp(x, dims: Sequence[int]):
   return torch.ops.tfl.transpose(x, dims)
 
 
+@register_decomp(torch.ops.aten.view.default)
+def _aten_view_decomp(x, shape: Sequence[int]):
+  return torch.ops.tfl.reshape(x, shape)
+
+
 @register_decomp(torch.ops.aten._softmax.default)
 def _aten__softmax_decomp(
     x, dim: int, half_to_float: bool  # pylint: disable=unused-argument

ai_edge_torch/odml_torch/experimental/torch_tfl/_lowerings.py

@@ -265,6 +265,22 @@ def _tfl_transpose_lowering(
   )
 
 
+@lower(torch.ops.tfl.reshape.default)
+def _tfl_reshape_lowering(
+    lctx: LoweringContext,
+    x: ir.Value,
+    shape: Sequence[int],
+) -> ir.Value:
+  constant_shape = lowering_utils.numpy_array_constant(
+      np.array(shape, dtype=np.int32)
+  )
+  return _ir_operation(
+      "tfl.reshape",
+      results=lowering_utils.node_meta_to_ir_types(lctx.node),
+      operands=[x, constant_shape],
+  )
+
+
 @lower(torch.ops.tfl.softmax.default)
 def _tfl_softmax_lowering(
     lctx: LoweringContext,

ai_edge_torch/odml_torch/experimental/torch_tfl/_ops.py

@@ -102,6 +102,13 @@ def tfl_transpose(input: torch.Tensor, perm: Sequence[int]) -> torch.Tensor:
   return torch.permute(input, perm).clone()
 
 
+@custom_op_with_fake("tfl::reshape")
+def tfl_reshape(input: torch.Tensor, shape: Sequence[int]) -> torch.Tensor:
+  assert torch.Size(shape).numel() == input.numel()
+
+  return input.view(shape).clone()
+
+
 @custom_op_with_fake("tfl::softmax")
 def tfl_softmax(x: torch.Tensor) -> torch.Tensor:
   return torch.nn.functional.softmax(x)

ai_edge_torch/odml_torch/experimental/torch_tfl/test/test_torch_tfl_impls.py

@@ -130,6 +130,9 @@ def _assert_export_and_close(
       ("aten_gelu_3", torch.ops.aten.gelu.default, (rnd(torch.float32, (1, 10)),), dict(approximate="tanh")),
       ("aten_permute_0", torch.ops.aten.permute.default, (rnd(torch.float32, (10, 10)), [0, 1],), dict()),
       ("aten_permute_1", torch.ops.aten.permute.default, (rnd(torch.float32, (1, 10)), [0, 1],), dict()),
+      ("aten_view_0", torch.ops.aten.view.default, (rnd(torch.float32, (10, 10)), [1, 100],), dict()),
+      ("aten_view_1", torch.ops.aten.view.default, (rnd(torch.float32, (1, 10)), [10, 1],), dict()),
+      ("aten_view_2", torch.ops.aten.view.default, (rnd(torch.float32, (10, 10)), [2, 5, 10],), dict()),
       ("aten__softmax_0", torch.ops.aten._softmax.default, (rnd(torch.float32, (10, 10)), -1, False), dict()),
       ("aten__softmax_1", torch.ops.aten._softmax.default, (rnd(torch.float32, (1, 10)), -1, False), dict()),
       ("aten__softmax_2", torch.ops.aten._softmax.default, (rnd(torch.float32, (10, 10)), 0, False), dict()),