Method
call_map
(
self,
f: torch.fx.GraphModule,
mapped_args: List[ProxyValue],
operands: List[ProxyValue],
meta: NodeMetadata,
)
Source from the content-addressed store, hash-verified
| 564 | ) |
| 565 | |
| 566 | def call_map( |
| 567 | self, |
| 568 | f: torch.fx.GraphModule, |
| 569 | mapped_args: List[ProxyValue], |
| 570 | operands: List[ProxyValue], |
| 571 | meta: NodeMetadata, |
| 572 | ) -> ProxyValue: |
| 573 | xs = _unstack_pytree([arg.data for arg in mapped_args])[0] |
| 574 | f_branch = self.call_submodule(f, tuple(xs + [arg.data for arg in operands])) |
| 575 | assert f_branch is not None |
| 576 | return self._fx( |
| 577 | "call_function", |
| 578 | torch.ops.higher_order.map_impl, |
| 579 | (f_branch.graph_module, mapped_args, operands), |
| 580 | {}, |
| 581 | meta, |
| 582 | ) |
| 583 | |
| 584 | def call_scan( |
| 585 | self, |