Home
last modified time | relevance | path

Searched refs:decomposition_table (Results 1 – 25 of 30) sorted by relevance

12

/aosp_15_r20/external/pytorch/torch/nn/utils/_expanded_weights/
H A Dexpanded_weights_impl.py7 from torch._decomp import decomposition_table
19 decomposition_table[aten.rnn_relu.input],
20 decomposition_table[aten.rnn_relu.data],
23 decomposition_table[aten.rnn_tanh.input],
24 decomposition_table[aten.rnn_tanh.data],
27 decomposition_table[aten.lstm.input],
28 decomposition_table[aten.lstm.data],
31 decomposition_table[aten.gru.input],
32 decomposition_table[aten.gru.data],
/aosp_15_r20/external/pytorch/torch/onnx/_internal/fx/passes/
H A Ddecomp.py25 decomposition_table: Mapping[torch._ops.OpOverload, Callable],
30 self.decomposition_table = decomposition_table
72 decomposition_table=self.decomposition_table,
H A Dfunctionalization.py122 decomposition_table={},
/aosp_15_r20/external/pytorch/torch/onnx/_internal/fx/
H A Ddecomposition_table.py89 decomposition_table: dict[torch._ops.OperatorBase, Callable] = {}
97 …for op_overload, decomp_fn in torch._decomp.decomposition_table.items(): # type: ignore[attr-defi…
107 decomposition_table[op_overload] = decomp_fn
115 decomposition_table[op_overload] = decomp_fn
116 return decomposition_table
/aosp_15_r20/external/pytorch/torch/onnx/_internal/exporter/
H A D_decomp.py86 decomposition_table: dict[torch._ops.OperatorBase, Callable] = {}
91 …for op_overload, decomp_fn in torch._decomp.decomposition_table.items(): # type: ignore[attr-defi…
98 decomposition_table[op_overload] = decomp_fn
100 return decomposition_table
/aosp_15_r20/external/pytorch/torchgen/decompositions/
H A Dgen_jit_decompositions.py5 from torch.jit._decompositions import decomposition_table
59 [scripted_func.code for scripted_func in decomposition_table.values()] # type: ignore[misc]
65 for schema, scripted_func in decomposition_table.items():
/aosp_15_r20/external/pytorch/torch/fx/experimental/
H A Dproxy_tensor.py144 decomposition_table: Optional[Mapping[OpOverload, Callable]]
148 CURRENT_DECOMPOSITION_TABLE = decomposition_table or {}
1417 decomposition_table: Optional[Mapping[OpOverload, Callable]] = None,
1424 self.decomposition_table = decomposition_table or {}
1463 with decompose(self.decomposition_table), self.mode:
1783 decomposition_table: Optional[Mapping[OpOverload, Callable]],
1793 self.decomposition_table: Dict[OpOverload, Callable] = dict(
1794 decomposition_table or {}
1796 self.decomposition_table.setdefault(
2023 stack.enter_context(decompose(self.decomposition_table))
[all …]
/aosp_15_r20/external/pytorch/torch/onnx/_internal/
H A D_exporter_legacy.py36 decomposition_table,
301 decomposition_table: dict[torch._ops.OpOverload, Callable]
333 self.decomposition_table = options.decomposition_table
361 self.decomposition_table = (
362decomposition_table.create_onnx_friendly_decomposition_table( # type: ignore[assignment]
1251 options.decomposition_table,
H A Donnxruntime.py44 import torch.onnx._internal.fx.decomposition_table
95 decomposition_table,
768 import torch.onnx._internal.fx.decomposition_table
804 …support_dict = torch.onnx._internal.fx.decomposition_table._create_onnx_supports_op_overload_table(
1166 decompositions=self._resolved_onnx_exporter_options.decomposition_table,
/aosp_15_r20/external/pytorch/torch/_decomp/
H A Ddecompositions_for_jvp.py12 decomposition_table = torch._decomp.decomposition_table variable
64 elif decomp in decomposition_table:
65 decomposition_table_used = decomposition_table
H A D__init__.py34 decomposition_table = global_decomposition_table["post_autograd"] variable
/aosp_15_r20/external/pytorch/test/
H A Dtest_decomp.py13 from torch._decomp import core_aten_decompositions, decomposition_table
59 for k in decomposition_table
734 cfunc = make_fx(func, decomposition_table=decomposition_table)
781 cfunc = make_fx(func, decomposition_table=decomposition_table)
828 func not in decomposition_table
853 decomposition = decomposition_table[func]
1282 ops_missing_decomp = overloads_wanting_decomp - decomposition_table.keys()
1304 for op in decomposition_table.keys()
H A Dtest_proxy_tensor.py15 from torch._decomp import decomposition_table
354 out = make_fx(f, decomposition_table=factory_func_decomp)(torch.ones(2))
376 … traced = make_fx(f, decomposition_table=decomp, tracing_mode=self.tracing_mode)(torch.rand(3))
593 fx_module = make_fx(fn, tracing_mode=self.tracing_mode, decomposition_table=None)(x)
603 …silu_decomp_table = {torch.ops.aten.silu.default: decomposition_table[torch.ops.aten.silu.default]}
607 decomposition_table=silu_decomp_table,
737 decomposed_fx = make_fx(f, decomposition_table={aten.addmm.default: addmm})(*inps)
752 traced = make_fx(f, decomposition_table={torch.ops.aten.t.default: nop})(torch.randn(5))
1688 …r = str(make_fx(f, tracing_mode="symbolic", decomposition_table=decomposition_table)(torch.empty(4…
/aosp_15_r20/external/pytorch/torch/_functorch/
H A Dvmap.py287 from torch._decomp import decomposition_table
290 if decomp in decomposition_table:
291 VMAP_DECOMPOSITIONS_LIB.impl(decomp, decomposition_table[decomp])
/aosp_15_r20/external/pytorch/torch/jit/
H A D_decompositions.py15 decomposition_table: Dict[str, torch.jit.ScriptFunction] = {} variable
73 registry = decomposition_table
/aosp_15_r20/external/pytorch/torch/_prims/
H A Dcontext.py129 func = torch._decomp.decomposition_table.get(orig_func, None)
133 func = torch._decomp.decomposition_table.get(default, None)
/aosp_15_r20/external/pytorch/torch/_dynamo/
H A Deval_frame.py1231 decomposition_table: Optional[
1311 if decomposition_table is not None:
1542 decomposition_table=decomposition_table,
/aosp_15_r20/external/pytorch/torch/_subclasses/
H A Dfake_tensor.py225 from torch._decomp import decomposition_table
231 return decomposition_table[func].__module__.startswith(
233 ) and decomposition_table[func].__name__ in dir(decompositions)
1930 from torch._decomp import decomposition_table
1933 if func in decomposition_table and (
1943 return decomposition_table[func](*args, **kwargs)
/aosp_15_r20/external/executorch/backends/transforms/
H A Ddecompose_sdpa.py37 decomposition_table=get_decompositions( # pyre-fixme[6]
/aosp_15_r20/external/pytorch/test/dynamo/
H A Dtest_debug_utils.py34 fx = make_fx(fn, decomposition_table=decomps)(x)
/aosp_15_r20/external/pytorch/torch/_functorch/_aot_autograd/
H A Ddispatch_and_compile_graph.py56 decomposition_table=aot_config.decompositions,
/aosp_15_r20/external/pytorch/torch/_export/
H A D__init__.py151 decomposition_table=decomp_table,
/aosp_15_r20/external/pytorch/test/inductor/
H A Dtest_flex_decoding.py590 gm = make_fx(attention, decomposition_table={})(query, key, value)
602 gm = make_fx(attention, decomposition_table=core_aten_decompositions())(
H A Dtest_flex_attention.py600 gm = make_fx(flex_attention, decomposition_table={})(query, key, value)
612 gm = make_fx(flex_attention, decomposition_table=core_aten_decompositions())(
/aosp_15_r20/external/executorch/exir/
H A Dtracer.py671 decomposition_table=(

12