| /aosp_15_r20/external/pytorch/torch/_higher_order_ops/ |
| H A D | associative_scan.py | 31 def wrap_combine_fn_flat(*args, combine_fn, spec, num_leaves): argument
76 def __call__(self, combine_fn, input, dim): argument
337 def associative_scan_op_dense(combine_fn, input, dim): argument
347 def associative_scan_proxy_mode(mode, combine_fn, input, dim): argument
352 def assoiciative_scan_fake_tensor_mode(mode, combine_fn, input, dim): argument
358 def associative_scan_functionalize(ctx, combine_fn, input, dim): argument
|
| /aosp_15_r20/external/pytorch/torch/_inductor/runtime/ |
| H A D | triton_helpers.py | 271 combine_fn, argument
339 def exclusive_scan_decoupled_lookback_64(scratch_base, block_value, index, combine_fn): argument
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/ |
| H A D | collective_combiner_utils.h | 48 const std::function<Status(absl::Span<HloInstruction* const>)>& combine_fn, in CombineInstructionsByKey()
|
| /aosp_15_r20/external/pytorch/torch/_inductor/codegen/ |
| H A D | triton_split_scan.py | 81 def scan(self, dtypes, combine_fn, values): argument
|
| /aosp_15_r20/external/pytorch/torch/_inductor/ |
| H A D | loop_body.py | 392 def bind_scan_shim(self, combine_fn): argument
|
| H A D | dependencies.py | 711 def scan(self, dtypes, combine_fn, values): argument
|
| H A D | ops_handler.py | 766 def scan(dtypes, combine_fn, values) -> Tuple[None, ...]: argument
804 def scan(dtypes, combine_fn, values): argument
|
| H A D | lowering.py | 5587 def combine_fn(a_tuple, b_tuple): function
5611 def combine_fn(a_tuple, b_tuple): function
|
| H A D | ir.py | 559 def combine_fn(a, b): function
588 def combine_fn(a, b): function
|
| /aosp_15_r20/external/pytorch/torch/_dynamo/variables/ |
| H A D | higher_order_ops.py | 948 def arg_extractor(combine_fn, input, dim): argument
|
| /aosp_15_r20/external/pytorch/test/inductor/ |
| H A D | test_cuda_repro.py | 1198 def combine_fn(i, j): function
|
| /aosp_15_r20/external/pytorch/test/functorch/ |
| H A D | test_control_flow.py | 85 def _fake_associative_scan(combine_fn, input, dim, reverse=False): argument
|