Repository: pytorch
Issue: Register a torch op for `functorch.experimental.control_flow.map` to lower
Labels: triaged,module: xla,oncall: pt2,module: functorch,module: pt2-dispatcher

### 🚀 The feature, motivation and pitch

We'd like to lower this `functorch.experimental.control_flow.map` and support in XLA. It's not registered as a `torch` op and gets unrolled. What would be the formal/best way to get support this?

Sample code,
```
import torch
from functorch.experimental import control_flow
from typing import List

def custom_backend(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
    print("custom backend called with FX graph:")
    gm.graph.print_tabular()
    return gm.forward

torch._dynamo.reset()

def f(i, x):
  return x + i

def map1(x):
  return control_flow.map(f, torch.ones(5), x)

opt_map1 = torch.compile(map1, backend=custom_backend)

args = torch.ones(1, requires_grad=True) # (torch.randn((10, )), )
opt_map1(args)
```
fx.graph:
```
custom backend called with FX graph:
opcode         name        target                                                   args                      kwargs
-------------  ----------  -------------------------------------------------------  ------------------------  --------
placeholder    l_x_        L_x_                                                     ()                        {}
call_function  ones        <built-in method ones of type object at 0x7f6243a6bc00>  (5,)                      {}
call_method    getitem     __getitem__                                              (ones, 0)                 {}
get_attr       map_body_0  map_body_0                                               ()                        {}
call_function  map_1       map                                                      (map_body_0, ones, l_x_)  {}
output         output      output    
```

Dumped IR (unrolled)
```
IR {
  %0 = f32[] prim::Constant(), location=__call__@_ops.py:516, xla_shape=f32[], value=1
  %1 = f32[] prim::Constant(), location=__call__@_ops.py:516, xla_shape=f32[], value=1
  %2 = f32[] aten::mul(%1, %0), location=__call__@_ops.py:516, xla_shape=f32[]
  %3 = f32[1]{0} xla::device_data(), location=<module>@torchxla_control_flow_map.py:20, xla_shape=f32[1]{0}, device=TPU:0
  %4 = f32[1]{0} aten::add(%3, %2), location=__call__@_ops.py:516, xla_shape=f32[1]{0}
  %5 = f32[] prim::Constant(), location=__call__@_ops.py:516, xla_shape=f32[], value=1
  %6 = f32[] prim::Constant(), location=__call__@_ops.py:516, xla_shape=f32[], value=1
  %7 = f32[] aten::mul(%6, %5), location=__call__@_ops.py:516, xla_shape=f32[]
  %8 = f32[1]{0} aten::add(%3, %7), location=__call__@_ops.py:516, xla_shape=f32[1]{0}
  %9 = f32[] prim::Constant(), location=__call__@_ops.py:516, xla_shape=f32[], value=1
  %10 = f32[] prim::Constant(), location=__call__@_ops.py:516, xla_shape=f32[], value=1
  %11 = f32[] aten::mul(%10, %9), location=__call__@_ops.py:516, xla_shape=f32[]
  %12 = f32[1]{0} aten::add(%3, %11), location=__call__@_ops.py:516, xla_shape=f32[1]{0}
  %13 = f32[] prim::Constant(), location=__call__@_ops.py:516, xla_shape=f32[], value=1
  %14 = f32[] prim::Constant(), location=__call__@_ops.py:516, xla_shape=f32[], value=1
  %15 = f32[] aten::mul(%14, %13), location=__call__@_ops.py:516, xla_shape=f32[]
  %16 = f32[1]{0} aten::add(%3, %15), location=__call__@_ops.py:516, xla_shape=f32[1]{0}
  %17 = f32[] prim::Constant(), location=__call__@_ops.py:516, xla_shape=f32[], value=1
  %18 = f32[] prim::Constant(), location=__call__@_ops.py:516, xla_shape=f32[], value=1
  %19 = f32[] aten::mul(%18, %17), location=__call__@_ops.py:516, xla_shape=f32[]
  %20 = f32[1]{0} aten::add(%3, %19), location=__call__@_ops.py:516, xla_shape=f32[1]{0}
  %21 = f32[5,1]{1,0} aten::stack(%20, %16, %12, %8, %4), location=_stack_pytree@_map.py:300, xla_shape=f32[5,1]{1,0}, dim=0, ROOT=0
}
```

cc @bdhirsh @ezyang @msaroufim @wconstab @anijain2305 @zou3519 @Chillee @samdow @kshitij12345 @janeyx99 @JackCaoG



### Alternatives

We can try to work with the fx.graph, but feels like more of a hack.

### Additional context

This is also related to the earlier request, https://github.com/pytorch/pytorch/issues/50688