[MLX]: Add flip op handler for mlx backend (pytorch#22173)

### Summary

Add flip op handler for the mlx backend. Implemented it with the flip
handler in the [MLX op
library](https://ml-explore.github.io/mlx/build/html/cpp/ops.html#_CPPv44flipRK5arrayRKNSt6vectorIiEE14StreamOrDevice).

Fixes pytorch#19609

### Test plan
Wrote a test model for Flip and wrote a FlipTest with it. Tested it
with:
```
python3 -m executorch.backends.mlx.test.run_all_tests --rebuild flip
```

I have an M3 pro so I was able to run the tests

Author: bharath-123

backends/mlx/ops.py

@@ -81,6 +81,7 @@
     ExpandDimsNode,
     Expm1Node,
     ExpNode,
+    FlipNode,
     FloatOrVid,
     FloorDivideIntNode,
     FloorDivideNode,
@@ -5453,6 +5454,28 @@ def emit_reverse(in_slot, out_slot):
     return output_slots
 
 
+@REGISTRY.register(target=[torch.ops.aten.flip.default])
+def _flip_handler(P: MLXProgramBuilder, n: Node) -> Slot:
+    args = P.args(n)
+    require_args(args, 2, 2, "aten.flip")
+    require_kwargs(P.kwargs(n), set(), "aten.flip")
+    x, dims = args
+
+    require_static_ints(dims, "dims", "aten.flip")
+
+    out = P.make_or_get_slot(n)
+
+    P.emit(
+        FlipNode(
+            x=P.slot_to_tid(x),
+            out=P.slot_to_tid(out),
+            axes=list(dims),
+        )
+    )
+
+    return out
+
+
 # ---------------------------------------------------------------------------
 # In-place variants for ops with bespoke functional handlers (clamp, pow,
 # activations). Each reuses its functional handler via a passthrough that

backends/mlx/runtime/MLXInterpreter.h

@@ -837,6 +837,11 @@ exec_transpose(const TransposeNode& n, ExecutionState& st, StreamOrDevice s) {
   st.set_tensor(n.out, transpose(st.const_tensor_ref(n.x), n.perm, s));
 }
 
+inline void exec_flip(const FlipNode& n, ExecutionState& st, StreamOrDevice s) {
+  std::vector<int> axes(n.axes.begin(), n.axes.end());
+  st.set_tensor(n.out, flip(st.const_tensor_ref(n.x), axes, s));
+}
+
 inline void
 exec_as_strided(const AsStridedNode& n, ExecutionState& st, StreamOrDevice s) {
   const auto& x = st.const_tensor_ref(n.x);
@@ -2162,6 +2167,9 @@ class Interpreter {
       case OpCode::TRANSPOSE:
         ops::exec_transpose(std::get<TransposeNode>(instr.node), st, s);
         break;
+      case OpCode::FLIP:
+        ops::exec_flip(std::get<FlipNode>(instr.node), st, s);
+        break;
       case OpCode::AS_STRIDED:
         ops::exec_as_strided(std::get<AsStridedNode>(instr.node), st, s);
         break;

backends/mlx/serialization/schema.fbs

@@ -1056,6 +1056,12 @@ table MetalKernelNode {
     init_value: float = null;
 }
 
+table FlipNode {
+    x: Tid (required);
+    out: Tid (required);
+    axes: [int32] (required);
+}
+
 // =============================================================================
 // Union of all op types
 // =============================================================================
@@ -1198,6 +1204,7 @@ union OpNode {
     RandomBitsNode,
     UpdateAndAttendNode,
     TruncNode,
+    FlipNode
     // BC: Add new op nodes here (append only)
 }
 

backends/mlx/test/test_ops.py

@@ -8587,3 +8587,41 @@ def compute_expected_outputs(self, model, test_inputs):
                 return model(*test_inputs)
         finally:
             REGISTRY.uninstall(key)
+
+
+class FlipModel(nn.Module):
+    def __init__(self, dims: List[int]):
+        super().__init__()
+        self.dims = dims
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.flip(x, self.dims)
+
+
+@register_test
+class FlipTest(OpTestCase):
+    name = "flip"
+
+    def __init__(self, shape: Tuple[int, ...], dims: List[int]):
+        self.shape = shape
+        self.dims = dims
+        dims_str = "_".join(str(d) for d in dims)
+        shape_str = "x".join(str(s) for s in shape)
+        self.name = f"flip_{shape_str}_dims{dims_str}"
+
+    @classmethod
+    def get_test_configs(cls) -> List["FlipTest"]:
+        return [
+            cls(shape=(4, 5), dims=[0]),
+            cls(shape=(4, 5), dims=[1]),
+            cls(shape=(4, 5), dims=[0, 1]),
+            cls(shape=(3, 4, 5), dims=[-1]),
+            cls(shape=(3, 4, 5), dims=[0, 2]),
+            cls(shape=(3, 4, 5), dims=[0, 1, 2]),
+        ]
+
+    def create_model(self) -> nn.Module:
+        return FlipModel(self.dims)
+
+    def create_inputs(self) -> Tuple[torch.Tensor, ...]:
+        return (torch.randn(self.shape),)