Searched refs:flat_args_with_path (Results 1 – 5 of 5) sorted by relevance
| /aosp_15_r20/external/pytorch/torch/_inductor/ |
| H A D | __init__.py | 84 flat_args_with_path, received_spec = pytree.tree_flatten_with_path(
91 x[1] if isinstance(x[1], torch.Tensor) else None for x in flat_args_with_path
|
| /aosp_15_r20/external/pytorch/torch/_export/ |
| H A D | utils.py | 221 input_placeholders: List[torch.fx.Node], flat_args_with_path, range_constraints argument
247 if len(flat_args_with_path) != len(input_placeholders):
256 for (key_path, arg), node in zip(flat_args_with_path, input_placeholders):
755 flat_args_with_path, _ = tree_flatten_with_path(combined_args)
763 for (arg_path, arg), user_input_name in zip(flat_args_with_path, user_input_names):
|
| /aosp_15_r20/external/pytorch/torch/export/ |
| H A D | exported_program.py | 830 flat_args_with_path, received_spec = pytree.tree_flatten_with_path(
833 self._check_input_constraints(flat_args_with_path)
834 flat_args = tuple(x[1] for x in flat_args_with_path)
1102 def _check_input_constraints(self, flat_args_with_path): argument
1112 input_placeholders, flat_args_with_path, self.range_constraints
|
| H A D | _unlift.py | 24 flat_args_with_path, received_spec = pytree.tree_flatten_with_path(args)
36 flat_args_with_path,
|
| H A D | unflatten.py | 444 flat_args_with_path, in_spec = pytree.tree_flatten_with_path(
447 flat_args = [x[1] for x in flat_args_with_path]
499 new_flat_args_with_path = flat_args_with_path # type: ignore[assignment]
|