mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

[XLA][NFC] Make interface of module loaders consistent.

Browse Source 
LoadModuleFromData has (data, format, config, ...) signature while FromFile has (path, config, format, ...). Change the latter so `format` becomes the second argument in both cases.

Since I'm touching this file:
* Use `std::string_view` and `absl::Status`
* Change `ovr_config` parameter to `const &`

PiperOrigin-RevId: 601304308
This commit is contained in:
Oleg Shyshkov 2024-01-24 19:16:09 -08:00 committed by jax authors
parent 5fb8b29dd9
commit fb80d2abcb
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
examples/jax_cpp/main.cc
View File

@ -58,8 +58,9 @@ int main(int argc, char** argv) {
std::function<void(xla::HloModuleConfig*)> config_modifier_hook =
[](xla::HloModuleConfig* config) { config->set_seed(42); };
std::unique_ptr<xla::HloModule> test_module =
LoadModuleFromFile(hlo_filename, xla::hlo_module_loader_details::Config(),
"txt", config_modifier_hook)
LoadModuleFromFile(hlo_filename, /*format=*/"txt",
xla::hlo_module_loader_details::Config(),
config_modifier_hook)
.value();
const xla::HloModuleProto test_module_proto = test_module->ToProto();