Home
last modified time | relevance | path

Searched refs:PrepareModuleInput (Results 1 – 7 of 7) sorted by relevance

/aosp_15_r20/external/pytorch/test/distributed/tensor/parallel/
H A Dtest_tp_style.py19 PrepareModuleInput,
199 prepare_inp_style = PrepareModuleInput(
226 prepare_inps_dimension_mismatch = PrepareModuleInput(
230 prepare_inps_short_dimension = PrepareModuleInput(
246 prepare_inps = PrepareModuleInput(
273 prepare_inps_simple = PrepareModuleInput(
301 prepare_inps_simple = PrepareModuleInput(
H A Dtest_parallelize_api.py10 PrepareModuleInput,
177 PrepareModuleInput(
/aosp_15_r20/external/pytorch/torch/distributed/tensor/parallel/
H A D__init__.py7 PrepareModuleInput,
H A Dstyle.py333 class PrepareModuleInput(ParallelStyle): class
/aosp_15_r20/external/pytorch/test/distributed/_tensor/
H A Dtest_dtensor_compile.py33 PrepareModuleInput,
693 "attn": PrepareModuleInput(
751 prepare_module_input = PrepareModuleInput(
/aosp_15_r20/external/pytorch/torch/testing/_internal/distributed/_tensor/
H A Dcommon_dtensor.py21 PrepareModuleInput,
231 layer_parallelize_plan["attention"] = PrepareModuleInput(
/aosp_15_r20/external/pytorch/docs/source/
H A Ddistributed.tensor.parallel.rst41 .. autoclass:: torch.distributed.tensor.parallel.PrepareModuleInput