Searched refs:HloInputOutputAliasConfig (Results 1 – 19 of 19) sorted by relevance
40 const HloInputOutputAliasConfig& config) { in expect_aliased()47 absl::optional<HloInputOutputAliasConfig::Alias> aliased_param = in expect_aliased()57 const HloInputOutputAliasConfig& config) { in expect_not_aliased()63 absl::optional<HloInputOutputAliasConfig::Alias> aliased_param = in expect_not_aliased()85 HloInputOutputAliasConfig config( in TEST_F()91 /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); in TEST_F()117 HloInputOutputAliasConfig config( in TEST_F()123 /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); in TEST_F()128 /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); in TEST_F()156 HloInputOutputAliasConfig config( in TEST_F()[all …]
21 bool HloInputOutputAliasConfig::OutputHasAlias( in OutputHasAlias()26 Status HloInputOutputAliasConfig::SetUpAlias(const ShapeIndex& output_index, in SetUpAlias()54 HloInputOutputAliasProto HloInputOutputAliasConfig::ToProto() const { in ToProto()83 StatusOr<HloInputOutputAliasConfig> HloInputOutputAliasConfig::CreateFromProto( in CreateFromProto()85 HloInputOutputAliasConfig result(output_shape); in CreateFromProto()105 string HloInputOutputAliasConfig::ToString() const { in ToString()119 HloInputOutputAliasConfig::AliasKind120 HloInputOutputAliasConfig::ParameterAliasKind( in ParameterAliasKind()133 absl::optional<ShapeIndex> HloInputOutputAliasConfig::GetAliasedOutput( in GetAliasedOutput()146 absl::optional<HloInputOutputAliasConfig::Alias>[all …]
33 class HloInputOutputAliasConfig {57 HloInputOutputAliasConfig() = default;59 explicit HloInputOutputAliasConfig(Shape output_shape) in HloInputOutputAliasConfig() function62 virtual ~HloInputOutputAliasConfig() = default;88 static StatusOr<HloInputOutputAliasConfig> CreateFromProto(132 const HloInputOutputAliasConfig& config);
41 HloInputOutputAliasConfig* alias_config) { in Build()78 HloInputOutputAliasConfig::AliasKind::kSystemAlias)); in Build()98 HloInputOutputAliasConfig* alias_config = in Run()
53 [&](const ShapeIndex&, const HloInputOutputAliasConfig::Alias&) { in AliasCount()60 config_ = HloInputOutputAliasConfig(output_shape); in BuildAliasConfig()69 HloInputOutputAliasConfig config_;
240 HloInputOutputAliasConfig& input_output_alias_config() { in input_output_alias_config()243 const HloInputOutputAliasConfig& input_output_alias_config() const { in input_output_alias_config()321 HloInputOutputAliasConfig input_output_alias_config_;
245 /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); in TEST_F()248 /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); in TEST_F()253 /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); in TEST_F()291 /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); in TEST_F()294 /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); in TEST_F()299 /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); in TEST_F()382 /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); in TEST_F()385 /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); in TEST_F()
65 HloInputOutputAliasConfig* alias_config);
1380 /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); in TEST_F()1383 /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); in TEST_F()1415 /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); in TEST_F()1418 /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); in TEST_F()1483 /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); in TEST_F()1525 /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); in TEST_F()1567 /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); in TEST_F()
64 input_output_alias_config_ = HloInputOutputAliasConfig( in AddComputationInternal()346 HloInputOutputAliasConfig::CreateFromProto( in CreateFromProto()
354 const HloInputOutputAliasConfig::Alias& alias) { in AddCopiesForAliasedInputOutputs()401 const HloInputOutputAliasConfig::Alias& alias) -> Status { in AddCopiesForAliasedInputOutputs()
179 [this](const HloInputOutputAliasConfig::Alias& aliased_input) in ComputeInputOutputAliasedBuffers()
1542 const HloInputOutputAliasConfig::Alias& alias) { in BuildColocatedBufferSets()
490 {}, 0, {}, HloInputOutputAliasConfig::kUserAlias)); in TEST_F()
236 const xla::HloInputOutputAliasConfig& input_output_alias = in DoWork()240 const xla::HloInputOutputAliasConfig::Alias& alias) -> Status { in DoWork()242 return alias.kind == xla::HloInputOutputAliasConfig::AliasKind::kUserAlias in DoWork()
216 const HloInputOutputAliasConfig& input_output_alias = in CreateResultShapedBuffer()
329 const HloInputOutputAliasConfig& input_output_alias = in ExecuteOnStream()
968 const HloInputOutputAliasConfig& config = module->input_output_alias_config(); in TEST_F()
392 HloInputOutputAliasConfig config(program_shape.result()); in PopulateInputOutputAlias()412 HloInputOutputAliasConfig::AliasKind::kUserAlias)); in PopulateInputOutputAlias()