Searched refs:custom_call_handler_ (Results 1 – 6 of 6) sorted by relevance
69 custom_call_handler_(custom_call_handler), in slice_dynamic_output_()82 DynamicDimensionInference::CustomCallInferenceHandler custom_call_handler_; variable
137 CustomCallInferenceHandler custom_call_handler_; variable
157 custom_call_handler_ = std::move(handler); in set_custom_call_handler()372 custom_call_handler_; variable
77 custom_call_handler_(std::move(custom_call_handler)) {} in DynamicDimensionInferenceVisitor()197 DynamicDimensionInference::CustomCallInferenceHandler custom_call_handler_; member in xla::DynamicDimensionInferenceVisitor261 if (custom_call_handler_) { in HandleCustomCall()262 return custom_call_handler_(hlo, parent_); in HandleCustomCall()1664 : module_(module), custom_call_handler_(std::move(custom_call_handler)) {} in DynamicDimensionInference()1669 custom_call_handler_); in AnalyzeDynamicDimensions()
2516 if (!custom_call_handler_) { in HandleCustomCall()2530 auto output, custom_call_handler_(custom_call, absl::MakeSpan(operands))); in HandleCustomCall()
1810 DynamicDimensionInference::Run(module, custom_call_handler_)); in Run()