Skip to content

Commit c1a1ef0

Browse files
Merge pull request #844 from mlcommons/python_upgrades
fix: frozen dict conversion
2 parents 5c4c07d + a687fa7 commit c1a1ef0

File tree

2 files changed

+4
-1
lines changed

2 files changed

+4
-1
lines changed

tests/modeldiffs/diff.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
from flax import jax_utils
2+
from flax.core import FrozenDict
23
import jax
34
import numpy as np
45
import torch
@@ -16,6 +17,8 @@ def torch2jax(jax_workload,
1617
jax_params, model_state = jax_workload.init_model_fn(jax.random.PRNGKey(0),
1718
**init_kwargs)
1819
pytorch_model, _ = pytorch_workload.init_model_fn([0], **init_kwargs)
20+
if isinstance(jax_params, dict):
21+
jax_params = FrozenDict(jax_params)
1922
jax_params = jax_utils.unreplicate(jax_params).unfreeze()
2023
if model_state is not None:
2124
model_state = jax_utils.unreplicate(model_state)

tests/test_traindiffs.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ def test_workload(self, workload):
5050
pyt_logs = '/tmp/pyt_log.pkl'
5151
try:
5252
run(
53-
f'XLA_PYTHON_CLIENT_ALLOCATOR=platform python3 -m tests.reference_algorithm_tests --workload={workload} --framework=jax --global_batch_size={GLOBAL_BATCH_SIZE} --log_file={jax_logs}'
53+
f'XLA_PYTHON_CLIENT_ALLOCATOR=platform python -m tests.reference_algorithm_tests --workload={workload} --framework=jax --global_batch_size={GLOBAL_BATCH_SIZE} --log_file={jax_logs}'
5454
f' --submission_path=tests/modeldiffs/vanilla_sgd_jax.py --identical=True --tuning_search_space=None --num_train_steps={NUM_TRAIN_STEPS}',
5555
shell=True,
5656
stdout=DEVNULL,

0 commit comments

Comments
 (0)