Home
last modified time | relevance | path

Searched defs:TensorBox (Results 1 – 13 of 13) sorted by relevance

/aosp_15_r20/external/pytorch/torch/_inductor/
H A Dmkldnn_lowerings.py71 x: TensorBox,
72 weight: TensorBox,
73 bias: TensorBox,
99 x: TensorBox,
100 other: TensorBox,
101 weight: TensorBox,
102 bias: TensorBox,
133 x: TensorBox,
134 other: TensorBox,
135 weight: TensorBox,
[all …]
H A Dmkldnn_ir.py31 x: "TensorBox",
32 weight: "TensorBox",
33 bias: "TensorBox",
187 x: "TensorBox",
188 weight: "TensorBox",
189 bias: "TensorBox",
273 x: "TensorBox",
274 weight: "TensorBox",
275 bias: "TensorBox",
349 x: "TensorBox",
[all …]
H A Djagged_lowerings.py28 offsets: TensorBox,
111 jagged_values: TensorBox,
112 jagged_offsets: List[TensorBox],
182 dense: TensorBox,
183 jagged_offsets: List[TensorBox],
255 dense: TensorBox,
256 jagged_offsets: List[TensorBox],
H A Dir.py2898 def __init__(self, view: Union[BaseView, "TensorBox"]):
5787 x: "TensorBox",
5788 weight: "TensorBox",
5789 bias: "TensorBox",
5943 x: "TensorBox",
5944 weight: "TensorBox",
5945 bias: "TensorBox",
6029 x: "TensorBox",
6030 weight: "TensorBox",
6031 bias: "TensorBox",
[all …]
H A Dlowering.py467 def inner(*inputs: List[TensorBox], alpha=None):
549 def inner(*inputs: List[List[TensorBox]], alpha=1):
628 def to_dtype(x: TensorBox, dtype: torch.dtype, copy=False):
640 def _convert_element_type(x: TensorBox, dtype: torch.dtype):
655 def to_dtype_bitcast(x: TensorBox, dtype: torch.dtype, *, copy=False):
676 def _view_dtype(x: TensorBox, dtype: torch.dtype):
684 def to_device(x: TensorBox, device: torch.device, *, copy=False):
692 def _device_put(x: TensorBox, device: torch.device):
1143 input: TensorBox,
1144 scales: TensorBox,
[all …]
H A Ddependencies.py621 input_node: "torch._inductor.ir.TensorBox",
H A Dscheduler.py576 def get_buf_bytes(buf: Optional[Union[ir.Buffer, ir.TensorBox]]) -> int:
/aosp_15_r20/external/pytorch/torch/_inductor/kernel/
H A Dmm_scaled.py246 mat_a: TensorBox,
247 mat_b: TensorBox,
248 scale_a: TensorBox,
249 scale_b: TensorBox,
250 bias: Optional[TensorBox] = None,
251 scale_result: Optional[TensorBox] = None,
H A Dconv.py381 x: TensorBox,
382 weight: TensorBox,
383 bias: Optional[TensorBox],
446 x: TensorBox,
447 weight: TensorBox,
448 bias: TensorBox,
H A Dflex_attention.py91 args: List[TensorBox],
/aosp_15_r20/external/pytorch/torch/_inductor/fx_passes/
H A Db2b_gemm.py446 args: List[TensorBox],
516 A: torch._inductor.ir.TensorBox,
517 B: torch._inductor.ir.TensorBox,
518 C: torch._inductor.ir.TensorBox,
/aosp_15_r20/external/pytorch/torch/_inductor/codegen/
H A Dsimd.py1170 numel: sympy.Expr, buffers: Iterable[Union[ir.Buffer, ir.TensorBox]]
H A Dwrapper.py1009 self, code: IndentedBuffer, graph_inputs: Dict[str, ir.TensorBox]