google / google/flax

Imagenet example does not work with newer flax versions

Open
#3,364 32 comments 1 reaction 0 assignees View on GitHub
Dominant language
Jupyter Notebook
Stars
7.3k
Forks
833
Avg merge
5h 11m
Merged PRs (30d)
5

Description

Using newer versions of Flax, there seems to be an outdate API call in the imagenet test.

First, when testing with the most recent stable release of jax, 0.4.16 and flax 0.7.4, there was an orbax issue (most likely because orbax has not had a stable release since end of july.

Full error message (on TPU):
```
Using ssh batch size of 4. Attempting to SSH into 1 nodes with a total of 4 workers.
SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
2023-09-21 23:18:03.034775: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2023-09-21 23:18:03.053484: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2023-09-21 23:18:03.063527: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2023-09-21 23:18:03.110393: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2023-09-21 23:18:03.582526: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2023-09-21 23:18:03.582586: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2023-09-21 23:18:03.582593: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
2023-09-21 23:18:03.602052: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2023-09-21 23:18:03.602115: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2023-09-21 23:18:03.602121: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
2023-09-21 23:18:03.622315: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2023-09-21 23:18:03.622377: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2023-09-21 23:18:03.622384: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
2023-09-21 23:18:03.657278: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2023-09-21 23:18:03.657346: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2023-09-21 23:18:03.657353: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
Traceback (most recent call last):
File "/home/alijafari/flax/examples/imagenet/imagenet_fake_data_benchmark.py", line 32, in
import train
File "/home/alijafari/flax/examples/imagenet/train.py", line 33, in
from flax.training import train_state
File "/home/alijafari/.local/lib/python3.10/site-packages/flax/training/train_state.py", line 19, in
import optax
File "/home/alijafari/.local/lib/python3.10/site-packages/optax/__init__.py", line 102, in
from optax._src.second_order import fisher_diag
File "/home/alijafari/.local/lib/python3.10/site-packages/optax/_src/second_order.py", line 46, in
v: jnp.DeviceArray,
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/deprecations.py", line 53, in getattr
raise AttributeError(f"module {module!r} has no attribute {name!r}")
AttributeError: module 'jax.numpy' has no attribute 'DeviceArray'
Traceback (most recent call last):
File "/home/alijafari/flax/examples/imagenet/imagenet_fake_data_benchmark.py", line 32, in
import train
File "/home/alijafari/flax/examples/imagenet/train.py", line 33, in
from flax.training import train_state
File "/home/alijafari/.local/lib/python3.10/site-packages/flax/training/train_state.py", line 19, in
import optax
File "/home/alijafari/.local/lib/python3.10/site-packages/optax/__init__.py", line 102, in
from optax._src.second_order import fisher_diag
File "/home/alijafari/.local/lib/python3.10/site-packages/optax/_src/second_order.py", line 46, in
v: jnp.DeviceArray,
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/deprecations.py", line 53, in getattr
raise AttributeError(f"module {module!r} has no attribute {name!r}")
AttributeError: module 'jax.numpy' has no attribute 'DeviceArray'
Traceback (most recent call last):
File "/home/alijafari/flax/examples/imagenet/imagenet_fake_data_benchmark.py", line 32, in
import train
File "/home/alijafari/flax/examples/imagenet/train.py", line 33, in
from flax.training import train_state
File "/home/alijafari/.local/lib/python3.10/site-packages/flax/training/train_state.py", line 19, in
import optax
File "/home/alijafari/.local/lib/python3.10/site-packages/optax/__init__.py", line 102, in
from optax._src.second_order import fisher_diag
File "/home/alijafari/.local/lib/python3.10/site-packages/optax/_src/second_order.py", line 46, in
v: jnp.DeviceArray,
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/deprecations.py", line 53, in getattr
raise AttributeError(f"module {module!r} has no attribute {name!r}")
AttributeError: module 'jax.numpy' has no attribute 'DeviceArray'
Traceback (most recent call last):
File "/home/alijafari/flax/examples/imagenet/imagenet_fake_data_benchmark.py", line 32, in
import train
File "/home/alijafari/flax/examples/imagenet/train.py", line 33, in
from flax.training import train_state
File "/home/alijafari/.local/lib/python3.10/site-packages/flax/training/train_state.py", line 19, in
import optax
File "/home/alijafari/.local/lib/python3.10/site-packages/optax/__init__.py", line 102, in
from optax._src.second_order import fisher_diag
File "/home/alijafari/.local/lib/python3.10/site-packages/optax/_src/second_order.py", line 46, in
v: jnp.DeviceArray,
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/deprecations.py", line 53, in getattr
raise AttributeError(f"module {module!r} has no attribute {name!r}")
AttributeError: module 'jax.numpy' has no attribute 'DeviceArray'
##### Command execution on worker 0 failed with exit status 1. Continuing.
##### Command execution on worker 1 failed with exit status 1. Continuing.
##### Command execution on worker 2 failed with exit status 1. Continuing.
##### Command execution on worker 3 failed with exit status 1. Continuing
```

We then tested a jax and flax stable release that allign with the most recent orbax release: jax 0.4.13 and flax 0.7.0. However, this ran into a new error:

How to reproduce:
```
pip install "jax[tpu]==0.4.13" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
git clone https://github.com/google/flax.git && cd flax/examples/imagenet && pip install -r requirements.txt && pip install flax==0.7.0
cd flax/examples/imagenet && JAX_PLATFORMS=tpu python3 imagenet_fake_data_benchmark.py
```

Full error message (on TPU):
```
Using ssh batch size of 4. Attempting to SSH into 1 nodes with a total of 4 workers.
SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
2023-09-22 00:02:14.684726: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2023-09-22 00:02:14.714459: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2023-09-22 00:02:14.746449: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2023-09-22 00:02:14.755915: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2023-09-22 00:02:15.228367: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2023-09-22 00:02:15.228429: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2023-09-22 00:02:15.228435: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
2023-09-22 00:02:15.260470: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2023-09-22 00:02:15.260531: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2023-09-22 00:02:15.260537: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
2023-09-22 00:02:15.304349: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2023-09-22 00:02:15.304411: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2023-09-22 00:02:15.304417: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
2023-09-22 00:02:15.305893: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2023-09-22 00:02:15.305955: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2023-09-22 00:02:15.305961: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
Running tests under Python 3.10.12: /usr/bin/python3
2023-09-22 00:02:16.346399: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcuda.so.1'; dlerror: libcuda.so.1: cannot open shared object file: No such file or directory
2023-09-22 00:02:16.346423: W tensorflow/compiler/xla/stream_executor/cuda/cuda_driver.cc:265] failed call to cuInit: UNKNOWN ERROR (303)
[ RUN ] ImagenetBenchmarkFakeData.test_fake_data
Running tests under Python 3.10.12: /usr/bin/python3
2023-09-22 00:02:16.368624: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcuda.so.1'; dlerror: libcuda.so.1: cannot open shared object file: No such file or directory
2023-09-22 00:02:16.368648: W tensorflow/compiler/xla/stream_executor/cuda/cuda_driver.cc:265] failed call to cuInit: UNKNOWN ERROR (303)
[ RUN ] ImagenetBenchmarkFakeData.test_fake_data
Running tests under Python 3.10.12: /usr/bin/python3
2023-09-22 00:02:16.429315: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcuda.so.1'; dlerror: libcuda.so.1: cannot open shared object file: No such file or directory
2023-09-22 00:02:16.429338: W tensorflow/compiler/xla/stream_executor/cuda/cuda_driver.cc:265] failed call to cuInit: UNKNOWN ERROR (303)
[ RUN ] ImagenetBenchmarkFakeData.test_fake_data
Running tests under Python 3.10.12: /usr/bin/python3
2023-09-22 00:02:16.438815: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcuda.so.1'; dlerror: libcuda.so.1: cannot open shared object file: No such file or directory
2023-09-22 00:02:16.438839: W tensorflow/compiler/xla/stream_executor/cuda/cuda_driver.cc:265] failed call to cuInit: UNKNOWN ERROR (303)
[ RUN ] ImagenetBenchmarkFakeData.test_fake_data
I0922 00:02:26.652038 140051359447040 dataset_info.py:358] Load dataset info from /home/alijafari/flax/.tfds/metadata/imagenet2012/5.1.0
I0922 00:02:26.654772 140051359447040 dataset_info.py:411] Field info.description from disk and from code do not match. Keeping the one from code.
I0922 00:02:26.655005 140051359447040 dataset_info.py:411] Field info.module_name from disk and from code do not match. Keeping the one from code.
I0922 00:02:26.655146 140051359447040 mocking.py:151] Metadata found for imagenet2012 at /home/alijafari/flax/.tfds/metadata/imagenet2012/5.1.0
I0922 00:02:26.655780 140051359447040 logging_logger.py:35] Constructing tf.data.Dataset imagenet2012 for split train[0:320291], from /home/alijafari/flax/.tfds/metadata/imagenet2012/5.1.0
W0922 00:02:26.781402 140051359447040 options.py:588] options.experimental_threading is deprecated. Use options.threading instead.
WARNING:tensorflow:From /home/alijafari/.local/lib/python3.10/site-packages/tensorflow/python/autograph/pyct/static_analysis/liveness.py:83: Analyzer.lamba_check (from tensorflow.python.autograph.pyct.static_analysis.liveness) is deprecated and will be removed after 2023-09-23.
Instructions for updating:
Lambda fuctions will be no more assumed to be used in the statement where they are used, or at least in the same block. https://github.com/tensorflow/tensorflow/issues/56089
W0922 00:02:26.893633 140051359447040 deprecation.py:350] From /home/alijafari/.local/lib/python3.10/site-packages/tensorflow/python/autograph/pyct/static_analysis/liveness.py:83: Analyzer.lamba_check (from tensorflow.python.autograph.pyct.static_analysis.liveness) is deprecated and will be removed after 2023-09-23.
Instructions for updating:
Lambda fuctions will be no more assumed to be used in the statement where they are used, or at least in the same block. https://github.com/tensorflow/tensorflow/issues/56089
I0922 00:02:27.190137 140051359447040 mocking.py:151] Metadata found for imagenet2012 at /home/alijafari/flax/.tfds/metadata/imagenet2012/5.1.0
I0922 00:02:27.190793 140051359447040 logging_logger.py:35] Constructing tf.data.Dataset imagenet2012 for split validation[0:12500], from /home/alijafari/flax/.tfds/metadata/imagenet2012/5.1.0
W0922 00:02:27.212210 140051359447040 options.py:588] options.experimental_threading is deprecated. Use options.threading instead.
/home/alijafari/.local/lib/python3.10/site-packages/optax/_src/schedule.py:391: DeprecationWarning: jax.numpy.DeviceArray is deprecated. Use jax.Array.
def schedule(step: jnp.DeviceArray) -> jnp.DeviceArray:
[ FAILED ] ImagenetBenchmarkFakeData.test_fake_data
[ FAILED ] ImagenetBenchmarkFakeData.test_fake_data
======================================================================
ERROR: test_fake_data (__main__.ImagenetBenchmarkFakeData)
ImagenetBenchmarkFakeData.test_fake_data
----------------------------------------------------------------------
Traceback (most recent call last):
File "/home/alijafari/flax/examples/imagenet/imagenet_fake_data_benchmark.py", line 66, in
absltest.main()
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/testing/absltest.py", line 2049, in main
_run_in_app(run_tests, args, kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/testing/absltest.py", line 2166, in _run_in_app
app.run(main=main_function)
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/app.py", line 312, in run
_run_main(main, args)
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/app.py", line 258, in _run_main
sys.exit(main(argv))
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/testing/absltest.py", line 2164, in main_function
function(argv, args, kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/testing/absltest.py", line 2568, in run_tests
result = _run_and_get_tests_result(
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/testing/absltest.py", line 2537, in _run_and_get_tests_result
test_program = unittest.TestProgram(*args, **kwargs)
File "/usr/lib/python3.10/unittest/main.py", line 101, in __init__
self.runTests()
File "/usr/lib/python3.10/unittest/main.py", line 271, in runTests
self.result = testRunner.run(self.test)
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/testing/_pretty_print_reporter.py", line 86, in run
return super(TextTestRunner, self).run(test)
File "/usr/lib/python3.10/unittest/runner.py", line 184, in run
test(result)
File "/usr/lib/python3.10/unittest/suite.py", line 84, in __call__
return self.run(*args, **kwds)
File "/usr/lib/python3.10/unittest/suite.py", line 122, in run
test(result)
File "/usr/lib/python3.10/unittest/suite.py", line 84, in __call__
return self.run(*args, **kwds)
File "/usr/lib/python3.10/unittest/suite.py", line 122, in run
test(result)
File "/usr/lib/python3.10/unittest/case.py", line 650, in __call__
return self.run(*args, **kwds)
File "/usr/lib/python3.10/unittest/case.py", line 591, in run
self._callTestMethod(testMethod)
File "/usr/lib/python3.10/unittest/case.py", line 549, in _callTestMethod
method()
File "/home/alijafari/flax/examples/imagenet/imagenet_fake_data_benchmark.py", line 50, in test_fake_data
train.train_and_evaluate(config, workdir)
File "/home/alijafari/flax/examples/imagenet/train.py", line 361, in train_and_evaluate
state = create_train_state(rng, config, model, image_size, learning_rate_fn)
File "/home/alijafari/flax/examples/imagenet/train.py", line 258, in create_train_state
params, batch_stats = initialized(rng, image_size, model)
File "/home/alijafari/flax/examples/imagenet/train.py", line 69, in initialized
variables = init({'params': key}, jnp.ones(input_shape, model.dtype))
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 250, in cache_miss
outs, out_flat, out_tree, args_flat, jaxpr = _python_pjit_helper(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 158, in _python_pjit_helper
args_flat, _, params, in_tree, out_tree, _ = infer_params_fn(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/api.py", line 306, in infer_params
return pjit.common_infer_params(pjit_info_args, *args, **kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 505, in common_infer_params
jaxpr, consts, canonicalized_out_shardings_flat = _pjit_jaxpr(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 971, in _pjit_jaxpr
jaxpr, final_consts, out_type = _create_pjit_jaxpr(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 345, in memoized_fun
ans = call(fun, *args)
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 924, in _create_pjit_jaxpr
jaxpr, global_out_avals, consts = pe.trace_to_jaxpr_dynamic(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2155, in trace_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2177, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers_)
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 188, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/home/alijafari/flax/examples/imagenet/train.py", line 67, in init
return model.init(*args)
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/flax/linen/module.py", line 1662, in init
_, v_out = self.init_with_output(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/flax/linen/module.py", line 1567, in init_with_output
return init_with_output(
File "/home/alijafari/.local/lib/python3.10/site-packages/flax/core/scope.py", line 960, in wrapper
raise ValueError('First argument passed to an init function should be a '
jax._src.traceback_util.UnfilteredStackTrace: ValueError: First argument passed to an init function should be a `jax.PRNGKey` or a dictionary mapping strings to `jax.PRNGKey`.

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/home/alijafari/flax/examples/imagenet/imagenet_fake_data_benchmark.py", line 50, in test_fake_data
train.train_and_evaluate(config, workdir)
File "/home/alijafari/flax/examples/imagenet/train.py", line 361, in train_and_evaluate
state = create_train_state(rng, config, model, image_size, learning_rate_fn)
File "/home/alijafari/flax/examples/imagenet/train.py", line 258, in create_train_state
params, batch_stats = initialized(rng, image_size, model)
File "/home/alijafari/flax/examples/imagenet/train.py", line 69, in initialized
variables = init({'params': key}, jnp.ones(input_shape, model.dtype))
File "/home/alijafari/flax/examples/imagenet/train.py", line 67, in init
return model.init(*args)
ValueError: First argument passed to an init function should be a `jax.PRNGKey` or a dictionary mapping strings to `jax.PRNGKey`.

======================================================================
ERROR: test_fake_data (__main__.ImagenetBenchmarkFakeData)
ImagenetBenchmarkFakeData.test_fake_data
----------------------------------------------------------------------
Traceback (most recent call last):
File "/home/alijafari/.local/lib/python3.10/site-packages/flax/testing/benchmark.py", line 150, in tearDown
self._report_benchmark_results()
File "/home/alijafari/.local/lib/python3.10/site-packages/flax/testing/benchmark.py", line 258, in _report_benchmark_results
raise ValueError('Unable to determine test name for reporting '
ValueError: Unable to determine test name for reporting benchmark results. Make sure you are using `self.report_*` methods.

----------------------------------------------------------------------
Ran 1 test in 10.909s

FAILED (errors=2)
I0922 00:02:28.195271 139886586542080 dataset_info.py:358] Load dataset info from /home/alijafari/flax/.tfds/metadata/imagenet2012/5.1.0
I0922 00:02:28.198028 139886586542080 dataset_info.py:411] Field info.description from disk and from code do not match. Keeping the one from code.
I0922 00:02:28.198243 139886586542080 dataset_info.py:411] Field info.module_name from disk and from code do not match. Keeping the one from code.
I0922 00:02:28.198381 139886586542080 mocking.py:151] Metadata found for imagenet2012 at /home/alijafari/flax/.tfds/metadata/imagenet2012/5.1.0
I0922 00:02:28.200115 139886586542080 logging_logger.py:35] Constructing tf.data.Dataset imagenet2012 for split train[960873:1281164], from /home/alijafari/flax/.tfds/metadata/imagenet2012/5.1.0
W0922 00:02:28.339387 139886586542080 options.py:588] options.experimental_threading is deprecated. Use options.threading instead.
WARNING:tensorflow:From /home/alijafari/.local/lib/python3.10/site-packages/tensorflow/python/autograph/pyct/static_analysis/liveness.py:83: Analyzer.lamba_check (from tensorflow.python.autograph.pyct.static_analysis.liveness) is deprecated and will be removed after 2023-09-23.
Instructions for updating:
Lambda fuctions will be no more assumed to be used in the statement where they are used, or at least in the same block. https://github.com/tensorflow/tensorflow/issues/56089
W0922 00:02:28.452126 139886586542080 deprecation.py:350] From /home/alijafari/.local/lib/python3.10/site-packages/tensorflow/python/autograph/pyct/static_analysis/liveness.py:83: Analyzer.lamba_check (from tensorflow.python.autograph.pyct.static_analysis.liveness) is deprecated and will be removed after 2023-09-23.
Instructions for updating:
Lambda fuctions will be no more assumed to be used in the statement where they are used, or at least in the same block. https://github.com/tensorflow/tensorflow/issues/56089
I0922 00:02:28.746239 139886586542080 mocking.py:151] Metadata found for imagenet2012 at /home/alijafari/flax/.tfds/metadata/imagenet2012/5.1.0
I0922 00:02:28.746915 139886586542080 logging_logger.py:35] Constructing tf.data.Dataset imagenet2012 for split validation[37500:50000], from /home/alijafari/flax/.tfds/metadata/imagenet2012/5.1.0
W0922 00:02:28.768107 139886586542080 options.py:588] options.experimental_threading is deprecated. Use options.threading instead.
/home/alijafari/.local/lib/python3.10/site-packages/optax/_src/schedule.py:391: DeprecationWarning: jax.numpy.DeviceArray is deprecated. Use jax.Array.
def schedule(step: jnp.DeviceArray) -> jnp.DeviceArray:
[ FAILED ] ImagenetBenchmarkFakeData.test_fake_data
[ FAILED ] ImagenetBenchmarkFakeData.test_fake_data
======================================================================
ERROR: test_fake_data (__main__.ImagenetBenchmarkFakeData)
ImagenetBenchmarkFakeData.test_fake_data
----------------------------------------------------------------------
Traceback (most recent call last):
File "/home/alijafari/flax/examples/imagenet/imagenet_fake_data_benchmark.py", line 66, in
absltest.main()
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/testing/absltest.py", line 2049, in main
_run_in_app(run_tests, args, kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/testing/absltest.py", line 2166, in _run_in_app
app.run(main=main_function)
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/app.py", line 312, in run
_run_main(main, args)
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/app.py", line 258, in _run_main
sys.exit(main(argv))
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/testing/absltest.py", line 2164, in main_function
function(argv, args, kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/testing/absltest.py", line 2568, in run_tests
result = _run_and_get_tests_result(
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/testing/absltest.py", line 2537, in _run_and_get_tests_result
test_program = unittest.TestProgram(*args, **kwargs)
File "/usr/lib/python3.10/unittest/main.py", line 101, in __init__
self.runTests()
File "/usr/lib/python3.10/unittest/main.py", line 271, in runTests
self.result = testRunner.run(self.test)
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/testing/_pretty_print_reporter.py", line 86, in run
return super(TextTestRunner, self).run(test)
File "/usr/lib/python3.10/unittest/runner.py", line 184, in run
test(result)
File "/usr/lib/python3.10/unittest/suite.py", line 84, in __call__
return self.run(*args, **kwds)
File "/usr/lib/python3.10/unittest/suite.py", line 122, in run
test(result)
File "/usr/lib/python3.10/unittest/suite.py", line 84, in __call__
return self.run(*args, **kwds)
File "/usr/lib/python3.10/unittest/suite.py", line 122, in run
test(result)
File "/usr/lib/python3.10/unittest/case.py", line 650, in __call__
return self.run(*args, **kwds)
File "/usr/lib/python3.10/unittest/case.py", line 591, in run
self._callTestMethod(testMethod)
File "/usr/lib/python3.10/unittest/case.py", line 549, in _callTestMethod
method()
File "/home/alijafari/flax/examples/imagenet/imagenet_fake_data_benchmark.py", line 50, in test_fake_data
train.train_and_evaluate(config, workdir)
File "/home/alijafari/flax/examples/imagenet/train.py", line 361, in train_and_evaluate
state = create_train_state(rng, config, model, image_size, learning_rate_fn)
File "/home/alijafari/flax/examples/imagenet/train.py", line 258, in create_train_state
params, batch_stats = initialized(rng, image_size, model)
File "/home/alijafari/flax/examples/imagenet/train.py", line 69, in initialized
variables = init({'params': key}, jnp.ones(input_shape, model.dtype))
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 250, in cache_miss
outs, out_flat, out_tree, args_flat, jaxpr = _python_pjit_helper(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 158, in _python_pjit_helper
args_flat, _, params, in_tree, out_tree, _ = infer_params_fn(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/api.py", line 306, in infer_params
return pjit.common_infer_params(pjit_info_args, *args, **kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 505, in common_infer_params
jaxpr, consts, canonicalized_out_shardings_flat = _pjit_jaxpr(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 971, in _pjit_jaxpr
jaxpr, final_consts, out_type = _create_pjit_jaxpr(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 345, in memoized_fun
ans = call(fun, *args)
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 924, in _create_pjit_jaxpr
jaxpr, global_out_avals, consts = pe.trace_to_jaxpr_dynamic(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2155, in trace_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2177, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers_)
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 188, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/home/alijafari/flax/examples/imagenet/train.py", line 67, in init
return model.init(*args)
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/flax/linen/module.py", line 1662, in init
_, v_out = self.init_with_output(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/flax/linen/module.py", line 1567, in init_with_output
return init_with_output(
File "/home/alijafari/.local/lib/python3.10/site-packages/flax/core/scope.py", line 960, in wrapper
raise ValueError('First argument passed to an init function should be a '
jax._src.traceback_util.UnfilteredStackTrace: ValueError: First argument passed to an init function should be a `jax.PRNGKey` or a dictionary mapping strings to `jax.PRNGKey`.

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/home/alijafari/flax/examples/imagenet/imagenet_fake_data_benchmark.py", line 50, in test_fake_data
train.train_and_evaluate(config, workdir)
File "/home/alijafari/flax/examples/imagenet/train.py", line 361, in train_and_evaluate
state = create_train_state(rng, config, model, image_size, learning_rate_fn)
File "/home/alijafari/flax/examples/imagenet/train.py", line 258, in create_train_state
params, batch_stats = initialized(rng, image_size, model)
File "/home/alijafari/flax/examples/imagenet/train.py", line 69, in initialized
variables = init({'params': key}, jnp.ones(input_shape, model.dtype))
File "/home/alijafari/flax/examples/imagenet/train.py", line 67, in init
return model.init(*args)
ValueError: First argument passed to an init function should be a `jax.PRNGKey` or a dictionary mapping strings to `jax.PRNGKey`.

======================================================================
ERROR: test_fake_data (__main__.ImagenetBenchmarkFakeData)
ImagenetBenchmarkFakeData.test_fake_data
----------------------------------------------------------------------
Traceback (most recent call last):
File "/home/alijafari/.local/lib/python3.10/site-packages/flax/testing/benchmark.py", line 150, in tearDown
self._report_benchmark_results()
File "/home/alijafari/.local/lib/python3.10/site-packages/flax/testing/benchmark.py", line 258, in _report_benchmark_results
raise ValueError('Unable to determine test name for reporting '
ValueError: Unable to determine test name for reporting benchmark results. Make sure you are using `self.report_*` methods.

----------------------------------------------------------------------
Ran 1 test in 12.547s

FAILED (errors=2)
I0922 00:02:28.993444 139703183525888 dataset_info.py:358] Load dataset info from /home/alijafari/flax/.tfds/metadata/imagenet2012/5.1.0
I0922 00:02:28.996214 139703183525888 dataset_info.py:411] Field info.description from disk and from code do not match. Keeping the one from code.
I0922 00:02:28.996427 139703183525888 dataset_info.py:411] Field info.module_name from disk and from code do not match. Keeping the one from code.
I0922 00:02:28.996560 139703183525888 mocking.py:151] Metadata found for imagenet2012 at /home/alijafari/flax/.tfds/metadata/imagenet2012/5.1.0
I0922 00:02:28.998337 139703183525888 logging_logger.py:35] Constructing tf.data.Dataset imagenet2012 for split train[320291:640582], from /home/alijafari/flax/.tfds/metadata/imagenet2012/5.1.0
W0922 00:02:29.142283 139703183525888 options.py:588] options.experimental_threading is deprecated. Use options.threading instead.
WARNING:tensorflow:From /home/alijafari/.local/lib/python3.10/site-packages/tensorflow/python/autograph/pyct/static_analysis/liveness.py:83: Analyzer.lamba_check (from tensorflow.python.autograph.pyct.static_analysis.liveness) is deprecated and will be removed after 2023-09-23.
Instructions for updating:
Lambda fuctions will be no more assumed to be used in the statement where they are used, or at least in the same block. https://github.com/tensorflow/tensorflow/issues/56089
W0922 00:02:29.256080 139703183525888 deprecation.py:350] From /home/alijafari/.local/lib/python3.10/site-packages/tensorflow/python/autograph/pyct/static_analysis/liveness.py:83: Analyzer.lamba_check (from tensorflow.python.autograph.pyct.static_analysis.liveness) is deprecated and will be removed after 2023-09-23.
Instructions for updating:
Lambda fuctions will be no more assumed to be used in the statement where they are used, or at least in the same block. https://github.com/tensorflow/tensorflow/issues/56089
I0922 00:02:29.553791 139703183525888 mocking.py:151] Metadata found for imagenet2012 at /home/alijafari/flax/.tfds/metadata/imagenet2012/5.1.0
I0922 00:02:29.554453 139703183525888 logging_logger.py:35] Constructing tf.data.Dataset imagenet2012 for split validation[12500:25000], from /home/alijafari/flax/.tfds/metadata/imagenet2012/5.1.0
I0922 00:02:29.559327 139777785620480 dataset_info.py:358] Load dataset info from /home/alijafari/flax/.tfds/metadata/imagenet2012/5.1.0
I0922 00:02:29.562028 139777785620480 dataset_info.py:411] Field info.description from disk and from code do not match. Keeping the one from code.
I0922 00:02:29.562225 139777785620480 dataset_info.py:411] Field info.module_name from disk and from code do not match. Keeping the one from code.
I0922 00:02:29.562349 139777785620480 mocking.py:151] Metadata found for imagenet2012 at /home/alijafari/flax/.tfds/metadata/imagenet2012/5.1.0
I0922 00:02:29.563956 139777785620480 logging_logger.py:35] Constructing tf.data.Dataset imagenet2012 for split train[640582:960873], from /home/alijafari/flax/.tfds/metadata/imagenet2012/5.1.0
W0922 00:02:29.575566 139703183525888 options.py:588] options.experimental_threading is deprecated. Use options.threading instead.
/home/alijafari/.local/lib/python3.10/site-packages/optax/_src/schedule.py:391: DeprecationWarning: jax.numpy.DeviceArray is deprecated. Use jax.Array.
def schedule(step: jnp.DeviceArray) -> jnp.DeviceArray:
[ FAILED ] ImagenetBenchmarkFakeData.test_fake_data
[ FAILED ] ImagenetBenchmarkFakeData.test_fake_data
======================================================================
ERROR: test_fake_data (__main__.ImagenetBenchmarkFakeData)
ImagenetBenchmarkFakeData.test_fake_data
----------------------------------------------------------------------
Traceback (most recent call last):
File "/home/alijafari/flax/examples/imagenet/imagenet_fake_data_benchmark.py", line 66, in
absltest.main()
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/testing/absltest.py", line 2049, in main
_run_in_app(run_tests, args, kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/testing/absltest.py", line 2166, in _run_in_app
app.run(main=main_function)
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/app.py", line 312, in run
_run_main(main, args)
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/app.py", line 258, in _run_main
sys.exit(main(argv))
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/testing/absltest.py", line 2164, in main_function
function(argv, args, kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/testing/absltest.py", line 2568, in run_tests
result = _run_and_get_tests_result(
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/testing/absltest.py", line 2537, in _run_and_get_tests_result
test_program = unittest.TestProgram(*args, **kwargs)
File "/usr/lib/python3.10/unittest/main.py", line 101, in __init__
self.runTests()
File "/usr/lib/python3.10/unittest/main.py", line 271, in runTests
self.result = testRunner.run(self.test)
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/testing/_pretty_print_reporter.py", line 86, in run
return super(TextTestRunner, self).run(test)
File "/usr/lib/python3.10/unittest/runner.py", line 184, in run
test(result)
File "/usr/lib/python3.10/unittest/suite.py", line 84, in __call__
return self.run(*args, **kwds)
File "/usr/lib/python3.10/unittest/suite.py", line 122, in run
test(result)
File "/usr/lib/python3.10/unittest/suite.py", line 84, in __call__
return self.run(*args, **kwds)
File "/usr/lib/python3.10/unittest/suite.py", line 122, in run
test(result)
File "/usr/lib/python3.10/unittest/case.py", line 650, in __call__
return self.run(*args, **kwds)
File "/usr/lib/python3.10/unittest/case.py", line 591, in run
self._callTestMethod(testMethod)
File "/usr/lib/python3.10/unittest/case.py", line 549, in _callTestMethod
method()
File "/home/alijafari/flax/examples/imagenet/imagenet_fake_data_benchmark.py", line 50, in test_fake_data
train.train_and_evaluate(config, workdir)
File "/home/alijafari/flax/examples/imagenet/train.py", line 361, in train_and_evaluate
state = create_train_state(rng, config, model, image_size, learning_rate_fn)
File "/home/alijafari/flax/examples/imagenet/train.py", line 258, in create_train_state
params, batch_stats = initialized(rng, image_size, model)
File "/home/alijafari/flax/examples/imagenet/train.py", line 69, in initialized
variables = init({'params': key}, jnp.ones(input_shape, model.dtype))
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 250, in cache_miss
outs, out_flat, out_tree, args_flat, jaxpr = _python_pjit_helper(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 158, in _python_pjit_helper
args_flat, _, params, in_tree, out_tree, _ = infer_params_fn(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/api.py", line 306, in infer_params
return pjit.common_infer_params(pjit_info_args, *args, **kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 505, in common_infer_params
jaxpr, consts, canonicalized_out_shardings_flat = _pjit_jaxpr(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 971, in _pjit_jaxpr
jaxpr, final_consts, out_type = _create_pjit_jaxpr(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 345, in memoized_fun
ans = call(fun, *args)
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 924, in _create_pjit_jaxpr
jaxpr, global_out_avals, consts = pe.trace_to_jaxpr_dynamic(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2155, in trace_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2177, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers_)
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 188, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/home/alijafari/flax/examples/imagenet/train.py", line 67, in init
return model.init(*args)
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/flax/linen/module.py", line 1662, in init
_, v_out = self.init_with_output(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/flax/linen/module.py", line 1567, in init_with_output
return init_with_output(
File "/home/alijafari/.local/lib/python3.10/site-packages/flax/core/scope.py", line 960, in wrapper
raise ValueError('First argument passed to an init function should be a '
jax._src.traceback_util.UnfilteredStackTrace: ValueError: First argument passed to an init function should be a `jax.PRNGKey` or a dictionary mapping strings to `jax.PRNGKey`.

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/home/alijafari/flax/examples/imagenet/imagenet_fake_data_benchmark.py", line 50, in test_fake_data
train.train_and_evaluate(config, workdir)
File "/home/alijafari/flax/examples/imagenet/train.py", line 361, in train_and_evaluate
state = create_train_state(rng, config, model, image_size, learning_rate_fn)
File "/home/alijafari/flax/examples/imagenet/train.py", line 258, in create_train_state
params, batch_stats = initialized(rng, image_size, model)
File "/home/alijafari/flax/examples/imagenet/train.py", line 69, in initialized
variables = init({'params': key}, jnp.ones(input_shape, model.dtype))
File "/home/alijafari/flax/examples/imagenet/train.py", line 67, in init
return model.init(*args)
ValueError: First argument passed to an init function should be a `jax.PRNGKey` or a dictionary mapping strings to `jax.PRNGKey`.

======================================================================
ERROR: test_fake_data (__main__.ImagenetBenchmarkFakeData)
ImagenetBenchmarkFakeData.test_fake_data
----------------------------------------------------------------------
Traceback (most recent call last):
File "/home/alijafari/.local/lib/python3.10/site-packages/flax/testing/benchmark.py", line 150, in tearDown
self._report_benchmark_results()
File "/home/alijafari/.local/lib/python3.10/site-packages/flax/testing/benchmark.py", line 258, in _report_benchmark_results
raise ValueError('Unable to determine test name for reporting '
ValueError: Unable to determine test name for reporting benchmark results. Make sure you are using `self.report_*` methods.

----------------------------------------------------------------------
Ran 1 test in 13.262s

FAILED (errors=2)
W0922 00:02:29.722676 139777785620480 options.py:588] options.experimental_threading is deprecated. Use options.threading instead.
WARNING:tensorflow:From /home/alijafari/.local/lib/python3.10/site-packages/tensorflow/python/autograph/pyct/static_analysis/liveness.py:83: Analyzer.lamba_check (from tensorflow.python.autograph.pyct.static_analysis.liveness) is deprecated and will be removed after 2023-09-23.
Instructions for updating:
Lambda fuctions will be no more assumed to be used in the statement where they are used, or at least in the same block. https://github.com/tensorflow/tensorflow/issues/56089
W0922 00:02:29.835337 139777785620480 deprecation.py:350] From /home/alijafari/.local/lib/python3.10/site-packages/tensorflow/python/autograph/pyct/static_analysis/liveness.py:83: Analyzer.lamba_check (from tensorflow.python.autograph.pyct.static_analysis.liveness) is deprecated and will be removed after 2023-09-23.
Instructions for updating:
Lambda fuctions will be no more assumed to be used in the statement where they are used, or at least in the same block. https://github.com/tensorflow/tensorflow/issues/56089
I0922 00:02:30.127933 139777785620480 mocking.py:151] Metadata found for imagenet2012 at /home/alijafari/flax/.tfds/metadata/imagenet2012/5.1.0
I0922 00:02:30.128622 139777785620480 logging_logger.py:35] Constructing tf.data.Dataset imagenet2012 for split validation[25000:37500], from /home/alijafari/flax/.tfds/metadata/imagenet2012/5.1.0
W0922 00:02:30.149866 139777785620480 options.py:588] options.experimental_threading is deprecated. Use options.threading instead.
/home/alijafari/.local/lib/python3.10/site-packages/optax/_src/schedule.py:391: DeprecationWarning: jax.numpy.DeviceArray is deprecated. Use jax.Array.
def schedule(step: jnp.DeviceArray) -> jnp.DeviceArray:
[ FAILED ] ImagenetBenchmarkFakeData.test_fake_data
[ FAILED ] ImagenetBenchmarkFakeData.test_fake_data
======================================================================
ERROR: test_fake_data (__main__.ImagenetBenchmarkFakeData)
ImagenetBenchmarkFakeData.test_fake_data
----------------------------------------------------------------------
Traceback (most recent call last):
File "/home/alijafari/flax/examples/imagenet/imagenet_fake_data_benchmark.py", line 66, in
absltest.main()
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/testing/absltest.py", line 2049, in main
_run_in_app(run_tests, args, kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/testing/absltest.py", line 2166, in _run_in_app
app.run(main=main_function)
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/app.py", line 312, in run
_run_main(main, args)
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/app.py", line 258, in _run_main
sys.exit(main(argv))
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/testing/absltest.py", line 2164, in main_function
function(argv, args, kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/testing/absltest.py", line 2568, in run_tests
result = _run_and_get_tests_result(
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/testing/absltest.py", line 2537, in _run_and_get_tests_result
test_program = unittest.TestProgram(*args, **kwargs)
File "/usr/lib/python3.10/unittest/main.py", line 101, in __init__
self.runTests()
File "/usr/lib/python3.10/unittest/main.py", line 271, in runTests
self.result = testRunner.run(self.test)
File "/home/alijafari/.local/lib/python3.10/site-packages/absl/testing/_pretty_print_reporter.py", line 86, in run
return super(TextTestRunner, self).run(test)
File "/usr/lib/python3.10/unittest/runner.py", line 184, in run
test(result)
File "/usr/lib/python3.10/unittest/suite.py", line 84, in __call__
return self.run(*args, **kwds)
File "/usr/lib/python3.10/unittest/suite.py", line 122, in run
test(result)
File "/usr/lib/python3.10/unittest/suite.py", line 84, in __call__
return self.run(*args, **kwds)
File "/usr/lib/python3.10/unittest/suite.py", line 122, in run
test(result)
File "/usr/lib/python3.10/unittest/case.py", line 650, in __call__
return self.run(*args, **kwds)
File "/usr/lib/python3.10/unittest/case.py", line 591, in run
self._callTestMethod(testMethod)
File "/usr/lib/python3.10/unittest/case.py", line 549, in _callTestMethod
method()
File "/home/alijafari/flax/examples/imagenet/imagenet_fake_data_benchmark.py", line 50, in test_fake_data
train.train_and_evaluate(config, workdir)
File "/home/alijafari/flax/examples/imagenet/train.py", line 361, in train_and_evaluate
state = create_train_state(rng, config, model, image_size, learning_rate_fn)
File "/home/alijafari/flax/examples/imagenet/train.py", line 258, in create_train_state
params, batch_stats = initialized(rng, image_size, model)
File "/home/alijafari/flax/examples/imagenet/train.py", line 69, in initialized
variables = init({'params': key}, jnp.ones(input_shape, model.dtype))
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 250, in cache_miss
outs, out_flat, out_tree, args_flat, jaxpr = _python_pjit_helper(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 158, in _python_pjit_helper
args_flat, _, params, in_tree, out_tree, _ = infer_params_fn(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/api.py", line 306, in infer_params
return pjit.common_infer_params(pjit_info_args, *args, **kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 505, in common_infer_params
jaxpr, consts, canonicalized_out_shardings_flat = _pjit_jaxpr(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 971, in _pjit_jaxpr
jaxpr, final_consts, out_type = _create_pjit_jaxpr(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 345, in memoized_fun
ans = call(fun, *args)
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 924, in _create_pjit_jaxpr
jaxpr, global_out_avals, consts = pe.trace_to_jaxpr_dynamic(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2155, in trace_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2177, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers_)
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 188, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/home/alijafari/flax/examples/imagenet/train.py", line 67, in init
return model.init(*args)
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/flax/linen/module.py", line 1662, in init
_, v_out = self.init_with_output(
File "/home/alijafari/.local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/alijafari/.local/lib/python3.10/site-packages/flax/linen/module.py", line 1567, in init_with_output
return init_with_output(
File "/home/alijafari/.local/lib/python3.10/site-packages/flax/core/scope.py", line 960, in wrapper
raise ValueError('First argument passed to an init function should be a '
jax._src.traceback_util.UnfilteredStackTrace: ValueError: First argument passed to an init function should be a `jax.PRNGKey` or a dictionary mapping strings to `jax.PRNGKey`.

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/home/alijafari/flax/examples/imagenet/imagenet_fake_data_benchmark.py", line 50, in test_fake_data
train.train_and_evaluate(config, workdir)
File "/home/alijafari/flax/examples/imagenet/train.py", line 361, in train_and_evaluate
state = create_train_state(rng, config, model, image_size, learning_rate_fn)
File "/home/alijafari/flax/examples/imagenet/train.py", line 258, in create_train_state
params, batch_stats = initialized(rng, image_size, model)
File "/home/alijafari/flax/examples/imagenet/train.py", line 69, in initialized
variables = init({'params': key}, jnp.ones(input_shape, model.dtype))
File "/home/alijafari/flax/examples/imagenet/train.py", line 67, in init
return model.init(*args)
ValueError: First argument passed to an init function should be a `jax.PRNGKey` or a dictionary mapping strings to `jax.PRNGKey`.

======================================================================
ERROR: test_fake_data (__main__.ImagenetBenchmarkFakeData)
ImagenetBenchmarkFakeData.test_fake_data
----------------------------------------------------------------------
Traceback (most recent call last):
File "/home/alijafari/.local/lib/python3.10/site-packages/flax/testing/benchmark.py", line 150, in tearDown
self._report_benchmark_results()
File "/home/alijafari/.local/lib/python3.10/site-packages/flax/testing/benchmark.py", line 258, in _report_benchmark_results
raise ValueError('Unable to determine test name for reporting '
ValueError: Unable to determine test name for reporting benchmark results. Make sure you are using `self.report_*` methods.

----------------------------------------------------------------------
Ran 1 test in 13.907s

FAILED (errors=2)
##### Command execution on worker 2 failed with exit status 1. Continuing.
##### Command execution on worker 1 failed with exit status 1. Continuing.
##### Command execution on worker 3 failed with exit status 1. Continuing.
##### Command execution on worker 0 failed with exit status 1. Continuing.
```

### System information
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 22.04
- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`: see above
- Python version: 3.10
- GPU/TPU model and memory: TPU v5e (v5e-16)

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.