text:
Flax segfaults immediately on a Cloud TPU with the LM1B example

The LM1B example segfaults when running with a Cloud TPU v3-8. I followed the linked instructions to start a new tf-2-4-1 instance and ran the example with a TPU v3-8 in europe-west4-a. I can reproduce this error with a fresh instance and Cloud TPU. I am using a TFRC allocation.
Problem you have encountered:
Immediate segfault upon calling p_train_step. Error seems to still occur even after removing the pmap wrapper on train_step.
What you expected to happen:
LM1B training should continue without a segmentation fault. Training proceeds as normal on a NVIDIA Volta V100 GPU and on CPU.
Logs, error messages, etc:
Starting a tf-2-4-1 instance and TPU:
$ export ZONE=europe-west4-a
$ gcloud compute instances create $USER-user-vm-0001 \
   --machine-type=n1-standard-16 \
   --image-project=ml-images \
   --image-family=tf-2-4-1 \
   --boot-disk-size=200GB \
   --scopes=cloud-platform \
   --zone=$ZONE
$ gcloud compute tpus create $USER-tpu-0001 \
      --zone=$ZONE \
      --network=default \
      --accelerator-type=v3-8 \
      --version=tpu_driver_nightly

Running the LM1B example
paras@paras-user-vm-0001:~/flax/examples/lm1b$ pip3 install -r requirements.txt
<...omitted...>
paras@paras-user-vm-0001:~/flax/examples/lm1b$ COLAB_TPU_ADDR=<TPU_IP_HERE>:8470 python3 main.py --workdir=/tmp/lm1b --config=configs/default.py
I0429 06:59:47.689807 140126107940672 main.py:53] JAX host: 0 / 1
I0429 06:59:47.690321 140126107940672 main.py:54] JAX local devices: [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1), TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0), TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1), TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1), TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0), TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]
I0429 06:59:47.690537 140126107940672 local.py:45] Setting task status: host_id: 0, host_count: 1
I0429 06:59:47.690827 140126107940672 local.py:51] Created artifact workdir of type ArtifactType.DIRECTORY and value /tmp/lm1b.
I0429 06:59:47.691211 140126107940672 train.py:395] Initializing dataset.
I0429 06:59:47.691451 140126107940672 dataset_builder.py:812] No config specified, defaulting to first: lm1b/plain_text
I0429 06:59:47.692169 140126107940672 dataset_info.py:361] Load dataset info from /home/paras/tensorflow_datasets/lm1b/plain_text/1.0.0
I0429 06:59:47.694560 140126107940672 dataset_builder.py:511] Constructing tf.data.Dataset for split ReadInstruction([_RelativeInstruction(splitname='train', from_=0, to=30301028, unit='abs', rounding='closest')]), from /home/paras/tensorflow_datasets/lm1b/plain_text/1.0.0
I0429 06:59:47.837791 140126107940672 dataset_builder.py:812] No config specified, defaulting to first: lm1b/plain_text
I0429 06:59:47.838539 140126107940672 dataset_info.py:361] Load dataset info from /home/paras/tensorflow_datasets/lm1b/plain_text/1.0.0
I0429 06:59:47.840322 140126107940672 dataset_builder.py:511] Constructing tf.data.Dataset for split ReadInstruction([_RelativeInstruction(splitname='test', from_=0, to=306688, unit='abs', rounding='closest')]), from /home/paras/tensorflow_datasets/lm1b/plain_text/1.0.0
I0429 06:59:47.882424 140126107940672 tokenizer.py:136] SentencePiece vocab not found, building one from data.
sentencepiece_trainer.cc(177) LOG(INFO) Running command: --input=/tmp/ds_chars21vqp5v3 --vocab_size=30000 --character_coverage=1.0 --model_prefix=/tmp/sp_tmpqq2v3zpd --model_type=unigram
sentencepiece_trainer.cc(77) LOG(INFO) Starts training with :
trainer_spec {
  input: /tmp/ds_chars21vqp5v3
  input_format:
  model_prefix: /tmp/sp_tmpqq2v3zpd
  model_type: UNIGRAM
  vocab_size: 30000
  self_test_sample_size: 0
  character_coverage: 1
  input_sentence_size: 0
  shuffle_input_sentence: 1
  seed_sentencepiece_size: 1000000
  shrinking_factor: 0.75
  max_sentence_length: 4192
  num_threads: 16
  num_sub_iterations: 2
  max_sentencepiece_length: 16
  split_by_unicode_script: 1
  split_by_number: 1
  split_by_whitespace: 1
  split_digits: 0
  treat_whitespace_as_suffix: 0
  required_chars:
  byte_fallback: 0
  vocabulary_output_piece_score: 1
  train_extremely_large_corpus: 0
  hard_vocab_limit: 1
  use_all_vocab: 0
  unk_id: 0
  bos_id: 1
  eos_id: 2
  pad_id: -1
  unk_piece: <unk>
  bos_piece: <s>
  eos_piece: </s>
  pad_piece: <pad>
  unk_surface:  ⁇
}
normalizer_spec {
  name: nmt_nfkc
  add_dummy_prefix: 1
  remove_extra_whitespaces: 1
  escape_whitespaces: 1
  normalization_rule_tsv:
}
denormalizer_spec {}
trainer_interface.cc(319) LOG(INFO) SentenceIterator is not specified. Using MultiFileSentenceIterator.
trainer_interface.cc(174) LOG(INFO) Loading corpus: /tmp/ds_chars21vqp5v3
trainer_interface.cc(375) LOG(INFO) Loaded all 72730 sentences
trainer_interface.cc(390) LOG(INFO) Adding meta_piece: <unk>
trainer_interface.cc(390) LOG(INFO) Adding meta_piece: <s>
trainer_interface.cc(390) LOG(INFO) Adding meta_piece: </s>
trainer_interface.cc(395) LOG(INFO) Normalizing sentences...
trainer_interface.cc(456) LOG(INFO) all chars count=9997380
trainer_interface.cc(477) LOG(INFO) Alphabet size=129
trainer_interface.cc(478) LOG(INFO) Final character coverage=1
trainer_interface.cc(510) LOG(INFO) Done! preprocessed 72730 sentences.
unigram_model_trainer.cc(138) LOG(INFO) Making suffix array...
unigram_model_trainer.cc(142) LOG(INFO) Extracting frequent sub strings...
unigram_model_trainer.cc(193) LOG(INFO) Initialized 91621 seed sentencepieces
trainer_interface.cc(516) LOG(INFO) Tokenizing input sentences with whitespace: 72730
trainer_interface.cc(526) LOG(INFO) Done! 60525
unigram_model_trainer.cc(488) LOG(INFO) Using 60525 sentences for EM training
unigram_model_trainer.cc(504) LOG(INFO) EM sub_iter=0 size=45201 obj=11.1288 num_tokens=125913 num_tokens/piece=2.78562
unigram_model_trainer.cc(504) LOG(INFO) EM sub_iter=1 size=35948 obj=9.16787 num_tokens=126732 num_tokens/piece=3.52543
unigram_model_trainer.cc(504) LOG(INFO) EM sub_iter=0 size=30802 obj=9.09253 num_tokens=127524 num_tokens/piece=4.14012
unigram_model_trainer.cc(504) LOG(INFO) EM sub_iter=1 size=30756 obj=9.07286 num_tokens=128028 num_tokens/piece=4.1627
trainer_interface.cc(604) LOG(INFO) Saving model: /tmp/sp_tmpqq2v3zpd.model
trainer_interface.cc(615) LOG(INFO) Saving vocabs: /tmp/sp_tmpqq2v3zpd.vocab
I0429 07:00:07.270487 140126107940672 tokenizer.py:106] copied /tmp/sp_tmpqq2v3zpd.model to /tmp/lm1b/sentencepiece_model
I0429 07:00:09.691338 140126107940672 train.py:420] Initializing model, optimizer, and step functions.
I0429 07:05:24.729107 140126107940672 checkpoints.py:202] Found no checkpoint files in /tmp/lm1b
I0429 07:05:24.738514 140117427595008 logging_writer.py:56] Hyperparameters: {'attention_dropout_rate': 0.1, 'checkpoint_every_steps': 10000, 'dataset_name': 'lm1b', 'dropout_rate': 0.1, 'emb_dim': 512, 'eval_dataset_name': 'lm1b', 'eval_every_steps': 1000, 'eval_per_device_batch_size': 32, 'eval_split': 'test', 'label_smoothing': 0.0, 'learning_rate': 0.05, 'logits_via_embedding': False, 'max_corpus_chars': 10000000, 'max_eval_target_length': 512, 'max_predict_length': 50, 'max_target_length': 128, 'mlp_dim': 2048, 'num_eval_steps': 2000, 'num_heads': 8, 'num_layers': 6, 'num_predict_steps': -1, 'num_train_steps': 500000, 'per_device_batch_size': 128, 'prompts': 'I love to ', 'qkv_dim': 512, 'restore_checkpoints': True, 'sampling_temperature': 0.6, 'sampling_top_k': 20, 'save_checkpoints': True, 'seed': 0, 'use_bfloat16': True, 'vocab_path': '/tmp/lm1b/sentencepiece_model', 'vocab_size': 30000, 'warmup_steps': 1000, 'weight_decay': 0.1}
I0429 07:05:25.165973 140126107940672 train.py:511] Starting training loop.
Fatal Python error: Segmentation fault

Thread 0x00007f6fbaffd700 (most recent call first):
  File "/usr/lib/python3.7/multiprocessing/pool.py", line 470 in _handle_results
  File "/usr/lib/python3.7/threading.py", line 865 in run
  File "/usr/lib/python3.7/threading.py", line 917 in _bootstrap_inner
  File "/usr/lib/python3.7/threading.py", line 885 in _bootstrap

Thread 0x00007f6fbb7fe700 (most recent call first):
  File "/usr/lib/python3.7/multiprocessing/pool.py", line 422 in _handle_tasks
  File "/usr/lib/python3.7/threading.py", line 865 in run
  File "/usr/lib/python3.7/threading.py", line 917 in _bootstrap_inner
  File "/usr/lib/python3.7/threading.py", line 885 in _bootstrap

Thread 0x00007f6fbbfff700 (most recent call first):
  File "/usr/lib/python3.7/multiprocessing/pool.py", line 413 in _handle_workers
  File "/usr/lib/python3.7/threading.py", line 865 in run
  File "/usr/lib/python3.7/threading.py", line 917 in _bootstrap_inner
  File "/usr/lib/python3.7/threading.py", line 885 in _bootstrap

Thread 0x00007f6fe0ff9700 (most recent call first):
  File "/usr/lib/python3.7/multiprocessing/pool.py", line 110 in worker
  File "/usr/lib/python3.7/threading.py", line 865 in run
  File "/usr/lib/python3.7/threading.py", line 917 in _bootstrap_inner
  File "/usr/lib/python3.7/threading.py", line 885 in _bootstrap

Thread 0x00007f6fe17fa700 (most recent call first):
  File "/usr/lib/python3.7/multiprocessing/pool.py", line 470 in _handle_results
  File "/usr/lib/python3.7/threading.py", line 865 in run
  File "/usr/lib/python3.7/threading.py", line 917 in _bootstrap_inner
  File "/usr/lib/python3.7/threading.py", line 885 in _bootstrap

Thread 0x00007f6fa27fc700 (most recent call first):
  File "/usr/lib/python3.7/multiprocessing/pool.py", line 422 in _handle_tasks
  File "/usr/lib/python3.7/threading.py", line 865 in run
  File "/usr/lib/python3.7/threading.py", line 917 in _bootstrap_inner
  File "/usr/lib/python3.7/threading.py", line 885 in _bootstrap

Thread 0x00007f6fa1ffb700 (most recent call first):
  File "/usr/lib/python3.7/multiprocessing/pool.py", line 413 in _handle_workers
  File "/usr/lib/python3.7/threading.py", line 865 in run
  File "/usr/lib/python3.7/threading.py", line 917 in _bootstrap_inner
  File "/usr/lib/python3.7/threading.py", line 885 in _bootstrap

Thread 0x00007f6fa17fa700 (most recent call first):
  File "/usr/lib/python3.7/multiprocessing/pool.py", line 110 in worker
  File "/usr/lib/python3.7/threading.py", line 865 in run
  File "/usr/lib/python3.7/threading.py", line 917 in _bootstrap_inner
  File "/usr/lib/python3.7/threading.py", line 885 in _bootstrap

Current thread 0x00007f71a6e33740 (most recent call first):
  File "/home/paras/.local/lib/python3.7/site-packages/jax/interpreters/xla.py", line 1085 in _value
  File "/home/paras/.local/lib/python3.7/site-packages/jax/interpreters/xla.py", line 1210 in <lambda>
  File "/home/paras/.local/lib/python3.7/site-packages/jax/interpreters/xla.py", line 1159 in item
  File "/home/paras/.local/lib/python3.7/site-packages/jax/core.py", line 220 in __init__
  File "/home/paras/.local/lib/python3.7/site-packages/jax/interpreters/partial_eval.py", line 974 in lit
  File "/home/paras/.local/lib/python3.7/site-packages/jax/interpreters/partial_eval.py", line 979 in <listcomp>
  File "/home/paras/.local/lib/python3.7/site-packages/jax/interpreters/partial_eval.py", line 979 in _inline_literals
  File "/home/paras/.local/lib/python3.7/site-packages/jax/interpreters/partial_eval.py", line 943 in to_jaxpr
  File "/home/paras/.local/lib/python3.7/site-packages/jax/interpreters/partial_eval.py", line 1190 in trace_to_subjaxpr_dynamic
  File "/home/paras/.local/lib/python3.7/site-packages/jax/interpreters/partial_eval.py", line 1209 in trace_to_jaxpr_final
  File "/home/paras/.local/lib/python3.7/site-packages/jax/interpreters/pxla.py", line 688 in parallel_callable
  File "/home/paras/.local/lib/python3.7/site-packages/jax/linear_util.py", line 260 in memoized_fun
  File "/home/paras/.local/lib/python3.7/site-packages/jax/interpreters/pxla.py", line 618 in xla_pmap_impl
  File "/home/paras/.local/lib/python3.7/site-packages/jax/core.py", line 600 in process_call
  File "/home/paras/.local/lib/python3.7/site-packages/jax/core.py", line 1464 in process
  File "/home/paras/.local/lib/python3.7/site-packages/jax/core.py", line 1393 in call_bind
  File "/home/paras/.local/lib/python3.7/site-packages/jax/core.py", line 1461 in bind
  File "/home/paras/.local/lib/python3.7/site-packages/jax/api.py", line 1571 in f_pmapped
  File "/home/paras/.local/lib/python3.7/site-packages/jax/_src/traceback_util.py", line 139 in reraise_with_filtered_traceback
  File "/home/paras/flax/examples/lm1b/train.py", line 529 in train_and_evaluate
  File "main.py", line 63 in main
  File "/usr/local/lib/python3.7/dist-packages/absl/app.py", line 251 in _run_main
  File "/usr/local/lib/python3.7/dist-packages/absl/app.py", line 303 in run
  File "main.py", line 67 in <module>
Segmentation fault

Changes added to support TPU in LM1B example:
diff --git a/examples/lm1b/main.py b/examples/lm1b/main.py
index abcb10c..1ec9570 100644
--- a/examples/lm1b/main.py
+++ b/examples/lm1b/main.py
@@ -25,6 +25,7 @@ from absl import logging
 from clu import platform
 import train
 import jax
+import jax.tools.colab_tpu
 from ml_collections import config_flags
 import tensorflow as tf

@@ -43,6 +44,8 @@ def main(argv):
   if len(argv) > 1:
     raise app.UsageError('Too many command-line arguments.')

+  jax.tools.colab_tpu.setup_tpu()
+
   # Hide any GPUs form TensorFlow. Otherwise TF might reserve memory and make
   # it unavailable to JAX.
   tf.config.experimental.set_visible_devices([], 'GPU')

Steps to reproduce:
Please follow above instructions on a fresh Cloud TPU instance

Including results of a pip freeze: https://gist.github.com/parasj/949585942a4f6b96553de0c0910696fa

GDB backtrace: https://gist.github.com/parasj/d02f48d7ee6cc0a17d1428af2bc6db22
Thread 1 "main.py" received signal SIGSEGV, Segmentation fault.
0x00007ffff6d08725 in PyArray_DescrConverter () from /usr/local/lib/python3.7/dist-packages/numpy/core/_multiarray_umath.cpython-37m-x86_64-linux-gnu.so
(gdb) bt
#0  0x00007ffff6d08725 in PyArray_DescrConverter () from /usr/local/lib/python3.7/dist-packages/numpy/core/_multiarray_umath.cpython-37m-x86_64-linux-gnu.so
#1  0x00007fff797d0f08 in xla::PrimitiveTypeToDtype(xla::PrimitiveType) () from /home/paras/.local/lib/python3.7/site-packages/jaxlib/tpu_client_extension.so
#2  0x00007fff797d42a8 in xla::LiteralToPython(std::shared_ptr<xla::Literal>) () from /home/paras/.local/lib/python3.7/site-packages/jaxlib/tpu_client_extension.so
#3  0x00007fff795bb8bc in pybind11::cpp_function::initialize<xla::pybind11_init_tpu_client_extension(pybind11::module_&)::{lambda(xla::PyTpuBuffer*)#9}, stream_executor::port::StatusOr<pybind11::object>, xla::PyTpuBuffer*, pybind11::name, pybind11::is_method, pybind11::sibling>(xla::pybind11_init_tpu_client_extension(pybind11::module_&)::{lambda(xla::PyTpuBuffer*)#9}&&, stream_executor::port::StatusOr<pybind11::object> (*)(xla::PyTpuBuffer*), pybind11::name const&, pybind11::is_method const&, pybind11::sibling const&)::{lambda(pybind11::detail::function_call&)#3}::_FUN(pybind11::detail::function_call) ()
   from /home/paras/.local/lib/python3.7/site-packages/jaxlib/tpu_client_extension.so
#4  0x00007fff795b82a1 in pybind11::cpp_function::dispatcher(_object*, _object*, _object*) () from /home/paras/.local/lib/python3.7/site-packages/jaxlib/tpu_client_extension.so
#5  0x00000000005d88a4 in _PyMethodDef_RawFastCallKeywords ()
#6  0x00000000005da200 in _PyObject_FastCallKeywords ()
#7  0x000000000054b3c1 in ?? ()
#8  0x00000000005524cd in _PyEval_EvalFrameDefault ()
#9  0x00000000005da536 in _PyFunction_FastCallDict ()
#10 0x00000000004d5e22 in ?? ()

@avital I can confirm the segfault will still occur with a batch size of 64 and 16.

Looking at your run again, why are you using Colab configuration options? I don't think you should need those if you're running on a Cloud TPU.

I updated main.py to call into jax.tools.colab_tpu.setup_tpu() to set up the TPU configuration. I changed main.py to set up the TPU following the instructions on https://github.com/google/jax/blob/master/cloud_tpu_colabs/README.md#running-jax-on-a-cloud-tpu-from-a-gce-vm and still get a segfault.
diff --git a/examples/lm1b/main.py b/examples/lm1b/main.py
index abcb10c..8dd6f2b 100644
--- a/examples/lm1b/main.py
+++ b/examples/lm1b/main.py
@@ -25,6 +25,7 @@ from absl import logging
 from clu import platform
 import train
 import jax
+from jax.config import config
 from ml_collections import config_flags
 import tensorflow as tf

@@ -43,10 +44,8 @@ def main(argv):
   if len(argv) > 1:
     raise app.UsageError('Too many command-line arguments.')

-  # Hide any GPUs form TensorFlow. Otherwise TF might reserve memory and make
-  # it unavailable to JAX.
-  tf.config.experimental.set_visible_devices([], 'GPU')
-
+  config.FLAGS.jax_xla_backend = "tpu_driver"
+  config.FLAGS.jax_backend_target = "grpc://10.114.237.58:8470"
   logging.info('JAX host: %d / %d', jax.host_id(), jax.host_count())
   logging.info('JAX local devices: %r', jax.local_devices())

Do other example (not LM1b) work?

I tested the MNIST example (with the same modifification as above to set the jax_backend_target configuration setting. The MNIST example works:
paras@paras-user-vm-0001:~/flax/examples/mnist$ python3 main.py --workdir=/tmp/mnist --config=configs/default.py
I0504 17:09:08.876252 140626453968704 main.py:52] JAX host: 0 / 1
I0504 17:09:08.876626 140626453968704 main.py:53] JAX local devices: [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1), TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0), TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1), TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1), TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0), TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]
I0504 17:09:08.876908 140626453968704 local.py:45] Setting task status: host_id: 0, host_count: 1
I0504 17:09:08.877268 140626453968704 local.py:51] Created artifact workdir of type ArtifactType.DIRECTORY and value /tmp/mnist.
I0504 17:09:08.912700 140626453968704 dataset_info.py:358] Load dataset info from /home/paras/tensorflow_datasets/mnist/3.0.1
I0504 17:09:08.916606 140626453968704 dataset_info.py:417] Field info.citation from disk and from code do not match. Keeping the one from code.
I0504 17:09:08.916880 140626453968704 dataset_info.py:417] Field info.splits from disk and from code do not match. Keeping the one from code.
I0504 17:09:08.917015 140626453968704 dataset_info.py:417] Field info.module_name from disk and from code do not match. Keeping the one from code.
I0504 17:09:08.917197 140626453968704 dataset_builder.py:363] Reusing dataset mnist (/home/paras/tensorflow_datasets/mnist/3.0.1)
I0504 17:09:08.917343 140626453968704 logging_logger.py:34] Constructing tf.data.Dataset mnist for split train, from /home/paras/tensorflow_datasets/mnist/3.0.1
I0504 17:09:12.616118 140626453968704 logging_logger.py:34] Constructing tf.data.Dataset mnist for split test, from /home/paras/tensorflow_datasets/mnist/3.0.1
I0504 17:09:33.299320 140626453968704 train.py:126] train epoch: 1, loss: 0.2721, accuracy: 91.72
I0504 17:09:34.413738 140626453968704 train.py:177] eval epoch: 1, loss: 0.0496, accuracy: 98.38
I0504 17:09:42.884375 140626453968704 train.py:126] train epoch: 2, loss: 0.0513, accuracy: 98.40
I0504 17:09:42.926508 140626453968704 train.py:177] eval epoch: 2, loss: 0.0447, accuracy: 98.51
I0504 17:09:51.480505 140626453968704 train.py:126] train epoch: 3, loss: 0.0374, accuracy: 98.81
I0504 17:09:51.521068 140626453968704 train.py:177] eval epoch: 3, loss: 0.0270, accuracy: 99.02
I0504 17:10:00.194692 140626453968704 train.py:126] train epoch: 4, loss: 0.0257, accuracy: 99.15
I0504 17:10:00.246079 140626453968704 train.py:177] eval epoch: 4, loss: 0.0276, accuracy: 99.13
I0504 17:10:08.823210 140626453968704 train.py:126] train epoch: 5, loss: 0.0227, accuracy: 99.27
I0504 17:10:08.876967 140626453968704 train.py:177] eval epoch: 5, loss: 0.0250, accuracy: 99.21
I0504 17:10:17.498520 140626453968704 train.py:126] train epoch: 6, loss: 0.0174, accuracy: 99.47
I0504 17:10:17.558862 140626453968704 train.py:177] eval epoch: 6, loss: 0.0280, accuracy: 99.04
I0504 17:10:26.081090 140626453968704 train.py:126] train epoch: 7, loss: 0.0147, accuracy: 99.52
I0504 17:10:26.137810 140626453968704 train.py:177] eval epoch: 7, loss: 0.0267, accuracy: 99.08
I0504 17:10:34.764275 140626453968704 train.py:126] train epoch: 8, loss: 0.0113, accuracy: 99.65
I0504 17:10:34.807989 140626453968704 train.py:177] eval epoch: 8, loss: 0.0282, accuracy: 99.12

@avital Wondering if you have any updates on this or ideas how to debug LM1B?
This is a GDB backtrace after the segfault: https://gist.github.com/parasj/d02f48d7ee6cc0a17d1428af2bc6db22
paras@paras-user-vm-0001:~/flax/examples/lm1b$ COLAB_TPU_ADDR=10.114.237.58:8470 gdb --args python3 main.py --workdir=/tmp/lm1b --config=configs/default.py
GNU gdb (Debian 8.2.1-2+b3) 8.2.1
Copyright (C) 2018 Free Software Foundation, Inc.
License GPLv3+: GNU GPL version 3 or later <http://gnu.org/licenses/gpl.html>
This is free software: you are free to change and redistribute it.
There is NO WARRANTY, to the extent permitted by law.
Type "show copying" and "show warranty" for details.
This GDB was configured as "x86_64-linux-gnu".
Type "show configuration" for configuration details.
For bug reporting instructions, please see:
<http://www.gnu.org/software/gdb/bugs/>.
Find the GDB manual and other documentation resources online at:
    <http://www.gnu.org/software/gdb/documentation/>.

For help, type "help".
Type "apropos word" to search for commands related to "word"...
Reading symbols from python3...(no debugging symbols found)...done.
(gdb) r
Starting program: /usr/bin/python3 main.py --workdir=/tmp/lm1b --config=configs/default.py
[Thread debugging using libthread_db enabled]
Using host libthread_db library "/lib/x86_64-linux-gnu/libthread_db.so.1".
[New Thread 0x7ffff448d700 (LWP 25581)]
[New Thread 0x7ffff3c8c700 (LWP 25582)]
[New Thread 0x7fffef48b700 (LWP 25583)]
[New Thread 0x7fffecc8a700 (LWP 25584)]
[New Thread 0x7fffec489700 (LWP 25585)]
[New Thread 0x7fffe7c88700 (LWP 25586)]
[New Thread 0x7fffe7487700 (LWP 25587)]
[New Thread 0x7fffe4c86700 (LWP 25588)]
[New Thread 0x7fffe0485700 (LWP 25589)]
[New Thread 0x7fffdfc84700 (LWP 25590)]
[New Thread 0x7fffdb483700 (LWP 25591)]
[New Thread 0x7fffdac82700 (LWP 25592)]
[New Thread 0x7fffd8481700 (LWP 25593)]
[New Thread 0x7fffd5c80700 (LWP 25594)]
[New Thread 0x7fffd347f700 (LWP 25595)]
[Thread 0x7fffd347f700 (LWP 25595) exited]
[Thread 0x7fffd5c80700 (LWP 25594) exited]
[Thread 0x7fffd8481700 (LWP 25593) exited]
[Thread 0x7fffdac82700 (LWP 25592) exited]
[Thread 0x7fffdb483700 (LWP 25591) exited]
[Thread 0x7fffdfc84700 (LWP 25590) exited]
[Thread 0x7fffe0485700 (LWP 25589) exited]
[Thread 0x7fffe4c86700 (LWP 25588) exited]
[Thread 0x7fffe7487700 (LWP 25587) exited]
[Thread 0x7fffe7c88700 (LWP 25586) exited]
[Thread 0x7fffec489700 (LWP 25585) exited]
[Thread 0x7fffecc8a700 (LWP 25584) exited]
[Thread 0x7fffef48b700 (LWP 25583) exited]
[Thread 0x7ffff3c8c700 (LWP 25582) exited]
[Thread 0x7ffff448d700 (LWP 25581) exited]
[Detaching after fork from child process 25604]
[New Thread 0x7fffd347f700 (LWP 25606)]
[New Thread 0x7fffd5c80700 (LWP 25607)]
[New Thread 0x7fffd8481700 (LWP 25608)]
[New Thread 0x7fffdac82700 (LWP 25609)]
[New Thread 0x7fff9ee34700 (LWP 25610)]
[New Thread 0x7fff9c633700 (LWP 25611)]
[New Thread 0x7fff99e32700 (LWP 25612)]
[New Thread 0x7fff97631700 (LWP 25613)]
[New Thread 0x7fff94e30700 (LWP 25614)]
[New Thread 0x7fff9262f700 (LWP 25615)]
[New Thread 0x7fff8fe2e700 (LWP 25616)]
[New Thread 0x7fff8d62d700 (LWP 25617)]
[New Thread 0x7fff8ae2c700 (LWP 25618)]
[New Thread 0x7fff8862b700 (LWP 25619)]
[New Thread 0x7fff85e2a700 (LWP 25620)]
[New Thread 0x7fff6b5aa700 (LWP 25643)]
[New Thread 0x7fff62da9700 (LWP 25644)]
[New Thread 0x7fff6ada9700 (LWP 25645)]
[New Thread 0x7fff6a5a8700 (LWP 25646)]
[New Thread 0x7fff69da7700 (LWP 25649)]
[New Thread 0x7fff695a6700 (LWP 25650)]
[New Thread 0x7fff68da5700 (LWP 25651)]
[New Thread 0x7fff63fff700 (LWP 25652)]
[New Thread 0x7fff637fe700 (LWP 25653)]
[New Thread 0x7fff625a8700 (LWP 25654)]
[New Thread 0x7fff61da7700 (LWP 25655)]
[New Thread 0x7fff615a6700 (LWP 25656)]
[New Thread 0x7fff60da5700 (LWP 25657)]
[New Thread 0x7fff43fff700 (LWP 25658)]
[New Thread 0x7fff437fe700 (LWP 25659)]
[New Thread 0x7fff42ffd700 (LWP 25660)]
[New Thread 0x7fff427fc700 (LWP 25661)]
[New Thread 0x7fff41ffb700 (LWP 25662)]
[New Thread 0x7fff417fa700 (LWP 25663)]
[New Thread 0x7fff40ff9700 (LWP 25664)]
[New Thread 0x7fff33fff700 (LWP 25665)]
[New Thread 0x7fff337fe700 (LWP 25666)]
[New Thread 0x7fff32ffd700 (LWP 25667)]
[New Thread 0x7fff327fc700 (LWP 25668)]
[New Thread 0x7fff31ffb700 (LWP 25669)]
[New Thread 0x7fff317fa700 (LWP 25670)]
[New Thread 0x7fff30ff9700 (LWP 25671)]
[New Thread 0x7fff17fff700 (LWP 25672)]
[New Thread 0x7fff177fe700 (LWP 25673)]
[New Thread 0x7fff16ffd700 (LWP 25674)]
I0429 07:57:41.558467 140737347581760 main.py:53] JAX host: 0 / 1
I0429 07:57:41.558822 140737347581760 main.py:54] JAX local devices: [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1), TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0), TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1), TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1), TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0), TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]
I0429 07:57:41.559057 140737347581760 local.py:45] Setting task status: host_id: 0, host_count: 1
I0429 07:57:41.559405 140737347581760 local.py:51] Created artifact workdir of type ArtifactType.DIRECTORY and value /tmp/lm1b.
I0429 07:57:41.559699 140737347581760 train.py:395] Initializing dataset.
I0429 07:57:41.559938 140737347581760 dataset_builder.py:812] No config specified, defaulting to first: lm1b/plain_text
I0429 07:57:41.560672 140737347581760 dataset_info.py:361] Load dataset info from /home/paras/tensorflow_datasets/lm1b/plain_text/1.0.0
I0429 07:57:41.563244 140737347581760 dataset_builder.py:511] Constructing tf.data.Dataset for split ReadInstruction([_RelativeInstruction(splitname='train', from_=0, to=30301028, unit='abs', rounding='closest')]), from /home/paras/tensorflow_datasets/lm1b/plain_text/1.0.0
[New Thread 0x7fff167fc700 (LWP 25675)]
[New Thread 0x7fff15ffb700 (LWP 25676)]
[New Thread 0x7fff157fa700 (LWP 25677)]
[New Thread 0x7fff14ff9700 (LWP 25678)]
[New Thread 0x7ffefbfff700 (LWP 25679)]
[New Thread 0x7ffefb7fe700 (LWP 25680)]
[New Thread 0x7ffefaffd700 (LWP 25681)]
[New Thread 0x7ffefa7fc700 (LWP 25682)]
[New Thread 0x7ffef9ffb700 (LWP 25683)]
[New Thread 0x7ffef97fa700 (LWP 25684)]
[New Thread 0x7ffef8ff9700 (LWP 25685)]
[New Thread 0x7ffed7fff700 (LWP 25686)]
[New Thread 0x7ffed77fe700 (LWP 25687)]
[New Thread 0x7ffed6ffd700 (LWP 25688)]
[New Thread 0x7ffed67fc700 (LWP 25689)]
[New Thread 0x7ffed5ffb700 (LWP 25690)]
[New Thread 0x7ffed57fa700 (LWP 25691)]
[New Thread 0x7ffed4ff9700 (LWP 25692)]
[New Thread 0x7ffebbfff700 (LWP 25693)]
[New Thread 0x7ffebb7fe700 (LWP 25694)]
[New Thread 0x7ffebaffd700 (LWP 25695)]
[New Thread 0x7ffeba7fc700 (LWP 25696)]
[New Thread 0x7ffeb9ffb700 (LWP 25697)]
[New Thread 0x7ffeb97fa700 (LWP 25698)]
[New Thread 0x7ffeb8ff9700 (LWP 25699)]
[New Thread 0x7ffe97fff700 (LWP 25700)]
[New Thread 0x7ffe977fe700 (LWP 25701)]
[New Thread 0x7ffe96ffd700 (LWP 25702)]
[New Thread 0x7ffe967fc700 (LWP 25703)]
[New Thread 0x7ffe95ffb700 (LWP 25704)]
[New Thread 0x7ffe957fa700 (LWP 25705)]
[New Thread 0x7ffe94ff9700 (LWP 25706)]
[New Thread 0x7ffe73fff700 (LWP 25707)]
I0429 07:57:41.729943 140737347581760 dataset_builder.py:812] No config specified, defaulting to first: lm1b/plain_text
I0429 07:57:41.730669 140737347581760 dataset_info.py:361] Load dataset info from /home/paras/tensorflow_datasets/lm1b/plain_text/1.0.0
I0429 07:57:41.732458 140737347581760 dataset_builder.py:511] Constructing tf.data.Dataset for split ReadInstruction([_RelativeInstruction(splitname='test', from_=0, to=306688, unit='abs', rounding='closest')]), from /home/paras/tensorflow_datasets/lm1b/plain_text/1.0.0
[New Thread 0x7ffe737fe700 (LWP 25716)]
[Thread 0x7ffe737fe700 (LWP 25716) exited]
I0429 07:57:44.220124 140737347581760 train.py:420] Initializing model, optimizer, and step functions.
[New Thread 0x7ffe737fe700 (LWP 25823)]
[New Thread 0x7ffe72ffd700 (LWP 25824)]
[New Thread 0x7ffe727fc700 (LWP 25825)]
[New Thread 0x7ffe71ffb700 (LWP 25826)]
[Thread 0x7ffe71ffb700 (LWP 25826) exited]
[Thread 0x7ffe72ffd700 (LWP 25824) exited]
[Thread 0x7ffe727fc700 (LWP 25825) exited]
[Thread 0x7ffe737fe700 (LWP 25823) exited]
I0429 08:03:00.243278 140737347581760 checkpoints.py:202] Found no checkpoint files in /tmp/lm1b
[New Thread 0x7ffe737fe700 (LWP 25827)]
[New Thread 0x7ffe71ffb700 (LWP 25828)]
[New Thread 0x7ffe727fc700 (LWP 25829)]
[New Thread 0x7ffe72ffd700 (LWP 25830)]
[New Thread 0x7ffe717fa700 (LWP 25831)]
[New Thread 0x7ffe70ff9700 (LWP 25832)]
[New Thread 0x7ffe53fff700 (LWP 25833)]
[New Thread 0x7ffe537fe700 (LWP 25834)]
I0429 08:03:00.255097 140730836182784 logging_writer.py:56] Hyperparameters: {'attention_dropout_rate': 0.1, 'checkpoint_every_steps': 10000, 'dataset_name': 'lm1b', 'dropout_rate': 0.1, 'emb_dim': 512, 'eval_dataset_name': 'lm1b', 'eval_every_steps': 1000, 'eval_per_device_batch_size': 32, 'eval_split': 'test', 'label_smoothing': 0.0, 'learning_rate': 0.05, 'logits_via_embedding': False, 'max_corpus_chars': 10000000, 'max_eval_target_length': 512, 'max_predict_length': 50, 'max_target_length': 128, 'mlp_dim': 2048, 'num_eval_steps': 2000, 'num_heads': 8, 'num_layers': 6, 'num_predict_steps': -1, 'num_train_steps': 500000, 'per_device_batch_size': 128, 'prompts': 'I love to ', 'qkv_dim': 512, 'restore_checkpoints': True, 'sampling_temperature': 0.6, 'sampling_top_k': 20, 'save_checkpoints': True, 'seed': 0, 'use_bfloat16': True, 'vocab_path': '/tmp/lm1b/sentencepiece_model', 'vocab_size': 30000, 'warmup_steps': 1000, 'weight_decay': 0.1}
I0429 08:03:00.663384 140737347581760 train.py:511] Starting training loop.
[New Thread 0x7ffe52ffd700 (LWP 25835)]
[New Thread 0x7ffe527fc700 (LWP 25836)]
[New Thread 0x7ffe51ffb700 (LWP 25837)]
[New Thread 0x7ffe517fa700 (LWP 25838)]
[New Thread 0x7ffe50ff9700 (LWP 25839)]
[New Thread 0x7ffe13fff700 (LWP 25840)]
[New Thread 0x7ffe137fe700 (LWP 25841)]
[New Thread 0x7ffe12ffd700 (LWP 25842)]
[New Thread 0x7ffe127fc700 (LWP 25843)]
[New Thread 0x7ffe11ffb700 (LWP 25844)]
[New Thread 0x7ffe117fa700 (LWP 25845)]
[New Thread 0x7ffe10ff9700 (LWP 25846)]
[New Thread 0x7ffdf3fff700 (LWP 25847)]
[New Thread 0x7ffdf37fe700 (LWP 25848)]
[New Thread 0x7ffdf2ffd700 (LWP 25849)]
[New Thread 0x7ffdf27fc700 (LWP 25850)]
[New Thread 0x7ffdf1ffb700 (LWP 25851)]
[New Thread 0x7ffdf17fa700 (LWP 25852)]
[New Thread 0x7ffdf0ff9700 (LWP 25853)]
[New Thread 0x7ffdd3fff700 (LWP 25854)]
[New Thread 0x7ffdd37fe700 (LWP 25855)]
[New Thread 0x7ffdd2ffd700 (LWP 25856)]
[New Thread 0x7ffdd27fc700 (LWP 25857)]
[New Thread 0x7ffdd1ffb700 (LWP 25858)]
[New Thread 0x7ffdd17fa700 (LWP 25859)]
[New Thread 0x7ffdd0ff9700 (LWP 25860)]
[New Thread 0x7ffdbbfff700 (LWP 25861)]
[New Thread 0x7ffdbb7fe700 (LWP 25862)]
[New Thread 0x7ffdbaffd700 (LWP 25863)]
[New Thread 0x7ffdba7fc700 (LWP 25864)]
[New Thread 0x7ffdb9ffb700 (LWP 25865)]
[New Thread 0x7ffdb97fa700 (LWP 25866)]
[New Thread 0x7ffdb8ff9700 (LWP 25867)]
[New Thread 0x7ffd8ffff700 (LWP 25868)]
[New Thread 0x7ffd97fff700 (LWP 25869)]
[New Thread 0x7ffd977fe700 (LWP 25870)]
[New Thread 0x7ffd96ffd700 (LWP 25871)]
[New Thread 0x7ffd967fc700 (LWP 25872)]
[New Thread 0x7ffd95ffb700 (LWP 25873)]
[New Thread 0x7ffd957fa700 (LWP 25874)]
[New Thread 0x7ffd94ff9700 (LWP 25875)]
[New Thread 0x7ffd8f7fe700 (LWP 25876)]
[New Thread 0x7ffd8effd700 (LWP 25877)]
[New Thread 0x7ffd8e7fc700 (LWP 25878)]
[New Thread 0x7ffd8dffb700 (LWP 25879)]
[New Thread 0x7ffd8d7fa700 (LWP 25880)]
[New Thread 0x7ffd8cff9700 (LWP 25881)]
[New Thread 0x7ffd57fff700 (LWP 25882)]
[New Thread 0x7ffd577fe700 (LWP 25883)]
[New Thread 0x7ffd56ffd700 (LWP 25884)]
[New Thread 0x7ffd567fc700 (LWP 25885)]
[New Thread 0x7ffd55ffb700 (LWP 25886)]
[New Thread 0x7ffd557fa700 (LWP 25887)]
[New Thread 0x7ffd54ff9700 (LWP 25888)]
[New Thread 0x7ffd33fff700 (LWP 25889)]
[New Thread 0x7ffd337fe700 (LWP 25890)]
[New Thread 0x7ffd32ffd700 (LWP 25891)]
[New Thread 0x7ffd327fc700 (LWP 25892)]
[New Thread 0x7ffd31ffb700 (LWP 25893)]
[New Thread 0x7ffd317fa700 (LWP 25894)]
[New Thread 0x7ffd30ff9700 (LWP 25895)]
[New Thread 0x7ffd23fff700 (LWP 25896)]
[New Thread 0x7ffd237fe700 (LWP 25897)]
[New Thread 0x7ffd22ffd700 (LWP 25898)]
[New Thread 0x7ffd227fc700 (LWP 25899)]
[New Thread 0x7ffd21ffb700 (LWP 25900)]
[New Thread 0x7ffd217fa700 (LWP 25901)]
[New Thread 0x7ffd20ff9700 (LWP 25902)]
[New Thread 0x7ffd207f8700 (LWP 25903)]
[New Thread 0x7ffd1fff7700 (LWP 25904)]
[New Thread 0x7ffd1f7f6700 (LWP 25905)]
[New Thread 0x7ffd1eff5700 (LWP 25906)]
[New Thread 0x7ffd1e7f4700 (LWP 25907)]
[New Thread 0x7ffd1dff3700 (LWP 25908)]

Thread 1 "main.py" received signal SIGSEGV, Segmentation fault.
0x00007ffff6d08725 in PyArray_DescrConverter () from /usr/local/lib/python3.7/dist-packages/numpy/core/_multiarray_umath.cpython-37m-x86_64-linux-gnu.so
(gdb) bt
#0  0x00007ffff6d08725 in PyArray_DescrConverter () from /usr/local/lib/python3.7/dist-packages/numpy/core/_multiarray_umath.cpython-37m-x86_64-linux-gnu.so
#1  0x00007fff797d0f08 in xla::PrimitiveTypeToDtype(xla::PrimitiveType) () from /home/paras/.local/lib/python3.7/site-packages/jaxlib/tpu_client_extension.so
#2  0x00007fff797d42a8 in xla::LiteralToPython(std::shared_ptr<xla::Literal>) () from /home/paras/.local/lib/python3.7/site-packages/jaxlib/tpu_client_extension.so
#3  0x00007fff795bb8bc in pybind11::cpp_function::initialize<xla::pybind11_init_tpu_client_extension(pybind11::module_&)::{lambda(xla::PyTpuBuffer*)#9}, stream_executor::port::StatusOr<pybind11::object>, xla::PyTpuBuffer*, pybind11::name, pybind11::is_method, pybind11::sibling>(xla::pybind11_init_tpu_client_extension(pybind11::module_&)::{lambda(xla::PyTpuBuffer*)#9}&&, stream_executor::port::StatusOr<pybind11::object> (*)(xla::PyTpuBuffer*), pybind11::name const&, pybind11::is_method const&, pybind11::sibling const&)::{lambda(pybind11::detail::function_call&)#3}::_FUN(pybind11::detail::function_call) ()
   from /home/paras/.local/lib/python3.7/site-packages/jaxlib/tpu_client_extension.so
#4  0x00007fff795b82a1 in pybind11::cpp_function::dispatcher(_object*, _object*, _object*) () from /home/paras/.local/lib/python3.7/site-packages/jaxlib/tpu_client_extension.so
#5  0x00000000005d88a4 in _PyMethodDef_RawFastCallKeywords ()
#6  0x00000000005da200 in _PyObject_FastCallKeywords ()
#7  0x000000000054b3c1 in ?? ()
#8  0x00000000005524cd in _PyEval_EvalFrameDefault ()
#9  0x00000000005da536 in _PyFunction_FastCallDict ()
#10 0x00000000004d5e22 in ?? ()
#11 0x00000000005a718f in PyObject_GetAttr ()
#12 0x000000000054e4e5 in _PyEval_EvalFrameDefault ()
#13 0x00000000005da536 in _PyFunction_FastCallDict ()
#14 0x000000000058f2e4 in ?? ()
#15 0x00000000004e22ec in PyNumber_Float ()
#16 0x00000000005d4335 in ?? ()
#17 0x00000000005d9bd7 in _PyObject_FastCallKeywords ()
#18 0x000000000055274b in _PyEval_EvalFrameDefault ()
#19 0x00000000005d91fc in _PyFunction_FastCallKeywords ()
#20 0x000000000054e7e0 in _PyEval_EvalFrameDefault ()
#21 0x00000000005da536 in _PyFunction_FastCallDict ()
#22 0x0000000000591533 in ?? ()
#23 0x00000000005d9c59 in _PyObject_FastCallKeywords ()
#24 0x000000000055274b in _PyEval_EvalFrameDefault ()
#25 0x000000000054c328 in _PyEval_EvalCodeWithName ()
#26 0x00000000005d94f2 in _PyFunction_FastCallKeywords ()
#27 0x000000000054e5ac in _PyEval_EvalFrameDefault ()
#28 0x000000000054c328 in _PyEval_EvalCodeWithName ()
#29 0x00000000005d94f2 in _PyFunction_FastCallKeywords ()
#30 0x000000000054e5ac in _PyEval_EvalFrameDefault ()
#31 0x000000000054c328 in _PyEval_EvalCodeWithName ()
#32 0x00000000005d94f2 in _PyFunction_FastCallKeywords ()
#33 0x000000000054e5ac in _PyEval_EvalFrameDefault ()
#34 0x000000000054c328 in _PyEval_EvalCodeWithName ()
#35 0x00000000005d94f2 in _PyFunction_FastCallKeywords ()
#36 0x000000000054e7e0 in _PyEval_EvalFrameDefault ()
#37 0x00000000005d91fc in _PyFunction_FastCallKeywords ()
#38 0x000000000054e5ac in _PyEval_EvalFrameDefault ()
#39 0x000000000054bcc2 in _PyEval_EvalCodeWithName ()
#40 0x00000000005d94f2 in _PyFunction_FastCallKeywords ()
#41 0x000000000054b1f0 in ?? ()
#42 0x000000000054f45e in _PyEval_EvalFrameDefault ()
#43 0x000000000054c328 in _PyEval_EvalCodeWithName ()
#44 0x00000000005da6fe in _PyFunction_FastCallDict ()
#45 0x000000000054f987 in _PyEval_EvalFrameDefault ()
#46 0x000000000054c328 in _PyEval_EvalCodeWithName ()
#47 0x00000000005da6fe in _PyFunction_FastCallDict ()
#48 0x000000000054f987 in _PyEval_EvalFrameDefault ()
#49 0x000000000054c328 in _PyEval_EvalCodeWithName ()
#50 0x00000000005da6fe in _PyFunction_FastCallDict ()
#51 0x000000000054f987 in _PyEval_EvalFrameDefault ()
#52 0x00000000005d91fc in _PyFunction_FastCallKeywords ()
#53 0x000000000054e7e0 in _PyEval_EvalFrameDefault ()
#54 0x00000000005d91fc in _PyFunction_FastCallKeywords ()
#55 0x000000000054e7e0 in _PyEval_EvalFrameDefault ()
#56 0x000000000054c328 in _PyEval_EvalCodeWithName ()
#57 0x00000000005da6fe in _PyFunction_FastCallDict ()
#58 0x000000000054f987 in _PyEval_EvalFrameDefault ()
#59 0x000000000054bcc2 in _PyEval_EvalCodeWithName ()
#60 0x00000000005da6fe in _PyFunction_FastCallDict ()
#61 0x00000000004d97e2 in ?? ()
#62 0x00000000005dc4f6 in PyObject_Call ()
#63 0x000000000054f987 in _PyEval_EvalFrameDefault ()
#64 0x000000000054c328 in _PyEval_EvalCodeWithName ()
#65 0x00000000005da6fe in _PyFunction_FastCallDict ()
#66 0x000000000054f987 in _PyEval_EvalFrameDefault ()
#67 0x000000000054c328 in _PyEval_EvalCodeWithName ()
#68 0x00000000005d94f2 in _PyFunction_FastCallKeywords ()
#69 0x000000000054b1f0 in ?? ()
#70 0x000000000054f45e in _PyEval_EvalFrameDefault ()
#71 0x000000000054c328 in _PyEval_EvalCodeWithName ()
#72 0x00000000005d94f2 in _PyFunction_FastCallKeywords ()
#73 0x000000000054b1f0 in ?? ()
#74 0x00000000005524cd in _PyEval_EvalFrameDefault ()
#75 0x00000000005d91fc in _PyFunction_FastCallKeywords ()
#76 0x000000000054e5ac in _PyEval_EvalFrameDefault ()
#77 0x00000000005d91fc in _PyFunction_FastCallKeywords ()
#78 0x000000000054e5ac in _PyEval_EvalFrameDefault ()
#79 0x000000000054bcc2 in _PyEval_EvalCodeWithName ()
#80 0x00000000005d94f2 in _PyFunction_FastCallKeywords ()
#81 0x000000000054b1f0 in ?? ()
#82 0x00000000005524cd in _PyEval_EvalFrameDefault ()
#83 0x000000000054bcc2 in _PyEval_EvalCodeWithName ()
#84 0x000000000054e0a3 in PyEval_EvalCode ()
#85 0x0000000000630ce2 in ?? ()
#86 0x0000000000630d97 in PyRun_FileExFlags ()
#87 0x00000000006319ff in PyRun_SimpleFileExFlags ()
#88 0x000000000065432e in ?? ()
#89 0x000000000065468e in _Py_UnixMain ()
#90 0x00007ffff79e609b in __libc_start_main (main=0x4bc560 <main>, argc=4, argv=0x7fffffffe518, init=<optimized out>, fini=<optimized out>, rtld_fini=<optimized out>, stack_end=0x7fffffffe508) at ../csu/libc-start.c:308
#91 0x00000000005e0e8a in _start ()

Honestly, debugging TPU segfaults isn't something I'm knowledgable about. @skye is the expert on this, she may have some insight. To help bisect the problem, could you try some of the other examples that use pmap (like lm1b does)?

Hi, lemme try to repro this myself before you spend too much time debugging. I'll use the same setup as in your initial post (thanks for the detailed information, it helps!).

Just wanted to let you know that I'm able to repro. I'll need to dig in to see where the issue is.

Thank you for looking into this @skye! Glad to hear the issue was clear.

We've released jaxlib 0.1.66, which should include this fix. Give it a shot?

@skye I had a chance to try it again and it seems to be working! I'm now getting a training loss. Thank you for looking into this and getting a fix pushed!

authors: avital parasj skye

repo: google/flax


text:
[nnx] jit uses graph_flatten/unflatten

What does this PR do?
Uses the improved graph_flatten and graph_unflatten in nnx.jit to correctly propagate updates to the Module structure.

Codecov Report
All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 59.45%. Comparing base (8220154) to head (536e507).

Additional details and impacted files
@@            Coverage Diff             @@
##             main    #3771      +/-   ##
==========================================
- Coverage   59.49%   59.45%   -0.04%
==========================================
  Files         101      101
  Lines       12623    12624       +1
==========================================
- Hits         7510     7506       -4
- Misses       5113     5118       +5

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

source: gharchive/pull-request

authors: cgarciae codecov-commenter

license_type: permissive

license_source: bigquery


text:
github.com/google/go-github/v30/github not found?

I'm not a big go user, though I wanted to build an app that uses the url github.com/google/go-github/v30/github and the result for that page is now not found. Can anyone explain this phenomenon?
https://github.com/rhysd/notes-cli/blob/4b2a04b3249a99667c2023d1aed42bbb3467d820/go.sum#L32-L33

GO111MODULE=on go get -u github.com/rhysd/notes-cli as a simple user I do not understand what modules are, or why it's important to me installing the application, causing this failure.

Already a problem that is discussed in https://github.com/golang/go/issues/30515 closing as a result.

source: gharchive/issue

repo: google/go-github

license_source: github-api


text:
feat: generate release notes support

supports the new feature github announced today: generate release notes (link)
while there are no official docs yet, one of the product managers posted screenshots of it on twitter: https://twitter.com/MylesBorins/status/1445102641214861318
This implements both the new /releases/generate-notes endpoint and the new generate_release_notes release option.

Thanks for the review @gmlewis! Just fixed the based on this tweet.

@Parker77 - do you have time to review this?


text:
Fix TrustedKeyKeys typo

In the godoc of TrustedIDKeyHashes, it should say "Not required if TrustedIDKeys is provided." instead of "Not required if TrustedKeyKeys is provided.".

Thanks for the fix!

authors: deeglaze msanft

repo: google/go-sev-guest


text:
check length of data in decodeName

When parsing malformed DNS, possible index out of range here.  Discovered this trying to parse an iodine session.

I am not interested in signing a license agreement so I can open a pull request.  Holy crap.

Sorry that didn't work out for you.  Unfortunately, to receive any code, we
need to have a CLA that allows us to use said code.
On Wed, Oct 19, 2016 at 3:15 PM, Hayden Parker notifications@github.com
wrote:

Closed #241 https://github.com/google/gopacket/pull/241.
—
You are receiving this because you are subscribed to this thread.
Reply to this email directly, view it on GitHub
https://github.com/google/gopacket/pull/241#event-829732404, or mute
the thread
https://github.com/notifications/unsubscribe-auth/ABJ9tko6YgoscmSTibbyBn7q0TBVDOnbks5q1oiCgaJpZM4KbbKi
.

@gconnell Is that really a legal requirement?  How are other open sourced projects able to operate without CLAs if thats true?  Is there any way I can just slap and MIT license on this diff so there is no doubt its free to use?  If I understand it correctly merging this would be receiving code from another 3-clause BSD licensed fork of the same codebase, so Google has 0 obligation to use it in any particular way.
I'm genuinely curious about why this is a requirement, not trying to be argumentative, so please answer my questions if you have the time.  But I'm not really interested in trying to get this merged, I'll just merge future changes into my fork and continue to use mine.

IANAL, but here's my understanding of things.  Hope it helps :)
When you've authored your code snippet, you've done a creative act.  If you've done it under company time of any sort, that grants copyright of the snippet to your employer.  If not, it grants one to you.  I believe the latter is the case, here.  By owning the copyright to the material you've sent, you have the basis to make demands on what's done with it.
Transfering that copyright could be difficult and hard to automate (I'm not sure exactly how it would work).  So instead, Google's CLA says "keep your copyright, just please tell us that we can use the code as we see fit."  You can read over https://cla.developers.google.com/about/google-individual if you want, but in short, it never removes your copyright from your code... it remains yours.  It DOES, however, grant Google the right to use that code as it wants, including redistributing it, etc.  For example, section 2 gives us a copyright license (https://en.wikipedia.org/wiki/Public_copyright_license).  It doesn't give us the copyright, but it gives us legal authority to use the code.
In short, we live in a world where, were we to accept code without this, you MIGHT have legal standing to come back later and make demands on how we use it.  And that's really scary.
As to the question of why it's not a requirement for pure open-source stuff, I don't know.  My guess is that, in those cases, you're not providing code to a legal entity (a corporation).  Instead, you're becoming part of a group, which collectively decides what to do with the code.  Again, IANAL and that's just a guess.  Note that as open source projects get larger, some of them DO start using CLAs, see for example Apache's at https://www.apache.org/licenses/icla.txt.
Unfortunately, I really can't accept non-CLA'd code.  But I've read over Google's CLA a few times, and I like what I see there.  We don't take ownership or copyright of code, we simply request that you allow us to use it like we want later on.  And I think that's a pretty good way to go about things :)

authors: gconnell hkparker

repo: google/gopacket


text:
panic with very big font

I'd like to show big font, e.g. show a very big letter in the whole screen.
so I change font size from 12 to 1024.
defaultFont, err := driver.CreateFont(gxfont.Default, 1024)

I got panic
panic: runtime error: slice bounds out of range
goroutine 1 [running, locked to thread]:
github.com/google/gxui/drivers/gl.(*glyphPage).add(...)
gopath/src/github.com/google/gxui/drivers/gl/glyph_page.go:138 +0x567
github.com/google/gxui/drivers/gl.(*glyphTable).get(...)
...
main.main()
gopath/src/github.com/google/gxui/samples/progress_bar/main.go:50 +0x2a
It is OK if I use 512 instead.
How big could a font be?

You are hitting the maximum glyph page size which is 512x512. This seemed like a reasonable limit.
I have to ask, is there a good reason to change this? :)

I try to show a rune(international word or char) to occupy the almost the whole screen.
One rune one screen.
My screen resolution is 1920x1080. 512 is OK but 768 looks better.
Maybe I could change the resolution or scale, but not sure.

The limitation comes from font file or gxui?
If possible, could we remove the limitation?

The limitation comes from font file or gxui?
GXUI. I've briefly described the font rendering process here.

If possible, could we remove the limitation?

I could make a glyph page expand to fit the size of the glyph when the page is empty. In your case of enormous glyphs, this would mean each glyph would get its own page (texture), and so every giant glyph would require a separate OpenGL draw call. This would mean performance would be dire if you were to try drawing giant strings, but that doesn't sound like an issue for what you're doing.

Thanks ben.
yes. I don't worry about the performance. I only show very few glyphs.
Please let me know if you have time to create the sample. :)

It would be more graceful to panic with a custom message--in this case, it's not immediately apparent that it's because the font size is too big.
A simple if statement would add virtually no overhead but it would prove helpful

authors: Annonomus-Penguin ben-clayton zhangpy

repo: google/gxui


text:
Error:Execution failed for task ':server:appengineEndpointsGetClientLibs'. >

Error:Execution failed for task ':server:appengineEndpointsGetClientLibs'.

There was an error running endpoints command get-client-lib: webapis-discovery.appspot.com

Obsolete after the 2018 version rewrite.

authors: JoseAlcerreca WanderingAway

repo: google/iosched


text:
lambda function argument leads to segfault; Partial alternative works

Description
Caveat: Unfortunately, I have not been able to produce minimal working example of below issue as I only encounter this  deep inside a jitted training loop containing several nested vmaps etc.
Define
   @partial(jit, static_argnames=['A_fun'])
   def fsai2(A_fun, G0):
       n = G0.shape[0]

       def _calc_G_i(i, G0_i):
           d_ii = jnp.dot(G0_i, A_fun(G0_i))**-0.5
           return d_ii

       G = vmap(_calc_G_i, (0, 0))(jnp.arange(n), G0)
       return G

then
     A_mvp = lambda b: A@b
     P = fsai2(A_mvp, jnp.eye(K.shape[0]))

gives
Segmentation fault (core dumped).
Similarly,
     from functools import partial
     A_mvp = partial(jnp.matmul, A)
     P = fsai2(A_mvp, jnp.eye(K.shape[0]))

gives
Segmentation fault (core dumped).
As does this:
     from jax.tree_util import Partial
     A_mvp = partial(jnp.matmul, A)
     P = fsai2(A_mvp, jnp.eye(K.shape[0]))

However, the last block of code with tree_util.Partial, works and gives expected results if I remove this line @partial(jit, static_argnames=['A_fun'])    decorator from the fsai2 function definition, i.e. this works:
   def fsai2(A_fun, G0):
       n = G0.shape[0]

       def _calc_G_i(i, G0_i):
           d_ii = jnp.dot(G0_i, A_fun(G0_i))**-0.5
           return d_ii

       G = vmap(_calc_G_i, (0, 0))(jnp.arange(n), G0)
       return G

  A_mvp = Partial(jnp.matmul, A)
  fsai2(A_mvp, jnp.eye(K.shape[0]))

I tried this last way as it was mentioned here https://github.com/google/jax/issues/1443 . Also relevant possibly https://github.com/google/jax/issues/5609.
What jax/jaxlib version are you using?
0.4.1 both
Which accelerator(s) are you using?
GPU
Additional system info
3.9.12 Python
NVIDIA GPU info
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.60.13    Driver Version: 525.60.13    CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA A100-SXM...  On   | 00000000:01:00.0 Off |                    0 |
| N/A   36C    P0    62W / 500W |      0MiB / 81920MiB |      0%   E. Process |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
|   1  NVIDIA A100-SXM...  On   | 00000000:41:00.0 Off |                    0 |
| N/A   36C    P0    60W / 500W |      0MiB / 81920MiB |      0%   E. Process |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
|   2  NVIDIA A100-SXM...  On   | 00000000:81:00.0 Off |                    0 |
| N/A   36C    P0    58W / 500W |      0MiB / 81920MiB |      0%   E. Process |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
|   3  NVIDIA A100-SXM...  On   | 00000000:C1:00.0 Off |                    0 |
| N/A   35C    P0    61W / 500W |      0MiB / 81920MiB |      0%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

Should've guessed, as so often, upgrading a newer version of JAX solved this.

Should've guessed, as so often, upgrading a newer version of JAX solved this.

repo: google/jax


text:
TPU deadlock

Hello,
I am trying to train reformer model using Trax and JAX. The training fails on Google Colab because of memory limitation. When I run it on google cloud server + TPU, it hangs on the "trax.supervised.Trainer".
The warning is as follows:
2020-08-26 17:46:37.421334: W external/org_tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/tpu_client.cc:601] TPU Execute is taking a long time. This might be due to a deadlock between multiple TPU cores or a very slow program.
The code is very straight forward:
import requests
import os

from jax.config import config
config.FLAGS.jax_xla_backend = "tpu_driver"
config.FLAGS.jax_backend_target = "grpc://" + "10.206.164.18"
print(config.FLAGS.jax_backend_target)

from tensorflow.compat.v1.io.gfile import GFile
import gin
import os
import jax
import trax
from trax.data import inputs

import numpy as np
import jax.numpy as jnp

from scipy.special import softmax

import sentencepiece as spm
from sentencepiece import SentencePieceProcessor

import random, glob, os

def fake_data():
  with open("vocab.txt",'w') as f:
    f.write("[MASK]\nL\nA\nG\nV\nE\nS\nI\nK\nR\nD\nT\nP\nN\nQ\nF\nY\nM\nH\nC\nW\nX\nU\nB\nZ\nO")

  if not os.path.exists('dataset'):
    os.makedirs('dataset')

  with open("dataset/train_0.txt",'w') as f:
    for i in range(50):
      f.write("M A F S A E D V L K E Y D R R R R M E A L L L S L Y Y P N D R K L L D Y K E W S P P R V Q V E C P K A P V E W N N P P S E K G L I V G H F S G I K Y K G E K A Q A S E V D V N K M C C W V S K F K D A M R R Y Q G I Q T C K I P G K V L S D L D A K I K A Y N L T V E G V E G F V R Y S R V T K Q H V A A F L K E L R H S K Q Y E N V N L I H Y I L T D K R V D I Q H L E K D L V K D F K A L V E S A H R M R Q G H M I N V K Y I L Y Q L L K K H G H G P D G P D I L T V K T G S K G V L Y D D S F R K I Y T D L G W K F T P L\n")
      f.write("M S I I G A T R L Q N D K S D T Y S A G P C Y A G G C S A F T P R G T C G K D W D L G E Q T C A S G F C T S Q P L C A R I K K T Q V C G L R Y S S K G K D P L V S A E W D S R G A P Y V R C T Y D A D L I D T Q A Q V D Q F V S M F G E S P S L A E R Y C M R G V K N T A G E L V S R V S S D A D P A G G W C R K W Y S A H R G P D Q D A A L G S F C I K N P G A A D C K C I N R A S D P V Y Q K V K T L H A Y P D Q C W Y V P C A A D V G E L K M G T Q R D T P T N C P T Q V C Q I V F N M L D D G S V T M D D V K N T I N C D F S K Y V P P P P P P K P T P P T P P T P P T P P T P P T P P T P P T P R P V H N R K V M F F V A G A V L V A I L I S T V R W\n")
      f.write("M A S N T V S A Q G G S N R P V R D F S N I Q D V A Q F L L F D P I W N E Q P G S I V P W K M N R E Q A L A E R Y P E L Q T S E P S E D Y S G P V E S L E L L P L E I K L D I M Q Y L S W E Q I S W C K H P W L W T R W Y K D N V V R V S A I T F E D F Q R E Y A F P E K I Q E I H F T D T R A E E I K A I L E T T P N V T R L V I R R I D D M N Y N T H G D L G L D D L E F L T H L M V E D A C G F T D F W A P S L T H L T I K N L D M H P R W F G P V M D G I K S M Q S T L K Y L Y I F E T Y G V N K P F V Q W C T D N I E T F Y C T N S Y R Y E N V P R P I Y V W V L F Q E D E W H G Y R V E D N K F H R R Y M Y S T I L H K R D T D W V E N N P L K T P A Q V E M Y K F L L R I S Q L N R D G T G Y E S D S D P E N E H F D D E S F S S G E E D S S D E D D P T W A P D S D D S D W E T E T E E E P S V A A R I L E K G K L T I T N L M K S L G F K P K P K K I Q S I D R Y F C S L D S N Y N S E D E D F E Y D S D S E D D D S D S E D D C\n")
      f.write("M Y Q A I N P C P Q S W Y G S P Q L E R E I V C K M S G A P H Y P N Y Y P V H P N A L G G A W F D T S L N A R S L T T T P S L T T C T P P S L A A C T P P T S L G M V D S P P H I N P P R R I G T L C F D F G S A K S P Q R C E C V A S D R P S T T S N T A P D T Y R L L I T N S K T R K N N Y G T C R L E P L T Y G I\n")
      f.write("M A R P L L G K T S S V R R R L E S L S A C S I F F F L R K F C Q K M A S L V F L N S P V Y Q M S N I L L T E R R Q V D R A M G G S D D D G V M V V A L S P S D F K T V L G S A L L A V E R D M V H V V P K Y L Q T P G I L H D M L V L L T P I F G E A L S V D M S G A T D V M V Q Q I A T A G F V D V D P L H S S V S W K D N V S C P V A L L A V S N A V R T M M G Q P C Q V T L I I D V G T Q N I L R D L V N L P V E M S G D L Q V M A Y T K D P L G K V P A V G V S V F D S G S V Q K G D A H S V G A P D G L V S F H T H P V S S A V E L N Y H A G W P S N V D M S S L L T M K N L M H V V V A E E G L W T M A R T L S M Q R L T K V L T D A E K D V M R A A A F N L F L P L N E L R V M G T K D S N N K S L K T Y F E V F E T F T I G A L M K H S G V T P T A F V D R R W L D N T I Y H M G F I P W G R D M R F V V E Y D L D G T N P F L N T V P T L M S V K R K A K I Q E M F D N M V S R M V T S\n")
      f.write("M N A K Y D T D Q G V G R M L F L G T I G L A V V V G G L M A Y G Y Y Y D G K T P S S G T S F H T A S P S F S S R Y R Y\n")
      f.write("M R Y T V L I A L Q G A L L L L L L I D D G Q G Q S P Y P Y P G M P C N S S R Q C G L G T C V H S R C A H C S S D G T L C S P E D P T M V W P C C P E S S C Q L V V G L P S L V N H Y N C L P N Q C T D S S Q C P G G F G C M T R R S K C E L C K A D G E A C N S P Y L D W R K D K E C C S G Y C H T E A R G L E G V C I D P K K I F C T P K N P W Q L A P Y P P S Y H Q P T T L R P P T S L Y D S W L M S G F L V K S T T A P S T Q E E E D D Y\n")
      f.write("M Q N P L P E V M S P E H D K R T T T P M S K E A N K F I R E L D K K P G D L A V V S D F V K R N T G K R L P I G K R S N L Y V R I C D L S G T I Y M G E T F I L E S W E E L Y L P E P T K M E V L G T L E S C C G I P P F P E W I V M V G E D Q C V Y A Y G D E E I L L F A Y S V K Q L V E E G I Q E T G I S Y K Y P D D I S D V D E E V L Q Q D E E I Q K I R K K T R E F V D K D A Q E F Q D F L N S L D A S L L S\n")
      f.write("M D S L N E V C Y E Q I K G T F Y K G L F G D F P L I V D K K T G C F N A T K L C V L G G K R F V D W N K T L R S K K L I Q Y Y E T R C D I K T E S L L Y E I K G D N N D E I T K Q I T G T Y L P K E F I L D I A S W I S V E F Y D K C N N I I I N Y F V N E Y K T M D K K T L Q S K I N E V E E K M Q K L L N E K E E E L Q E K N D K I D E L I L F S K R M E E D R K K D R E M M I K Q E K M L R E L G I H L E D V S S Q N N E L I E K V D E Q V E Q N A V L N F K I D N I Q N K L E I A V E D R A P Q P K Q N L K R E R F I L L K R N D D Y Y P Y Y T I R A Q D I N A R S A L K R Q K N L Y N E V S V L L D L T C H P N S K T L Y V R V K D E L K Q K G V V F N L C K V S I S N S K I N E E E L I K A M E T I N D E K R D V\n")

    with open("dataset/train_1.txt",'w') as f:
      for i in range(50):
        f.write("M A F S A E D V L K E Y D R R R R M E A L L L S L Y Y P N D R K L L D Y K E W S P P R V Q V E C P K A P V E W N N P P S E K G L I V G H F S G I K Y K G E K A Q A S E V D V N K M C C W V S K F K D A M R R Y Q G I Q T C K I P G K V L S D L D A K I K A Y N L T V E G V E G F V R Y S R V T K Q H V A A F L K E L R H S K Q Y E N V N L I H Y I L T D K R V D I Q H L E K D L V K D F K A L V E S A H R M R Q G H M I N V K Y I L Y Q L L K K H G H G P D G P D I L T V K T G S K G V L Y D D S F R K I Y T D L G W K F T P L\n")
        f.write("M S I I G A T R L Q N D K S D T Y S A G P C Y A G G C S A F T P R G T C G K D W D L G E Q T C A S G F C T S Q P L C A R I K K T Q V C G L R Y S S K G K D P L V S A E W D S R G A P Y V R C T Y D A D L I D T Q A Q V D Q F V S M F G E S P S L A E R Y C M R G V K N T A G E L V S R V S S D A D P A G G W C R K W Y S A H R G P D Q D A A L G S F C I K N P G A A D C K C I N R A S D P V Y Q K V K T L H A Y P D Q C W Y V P C A A D V G E L K M G T Q R D T P T N C P T Q V C Q I V F N M L D D G S V T M D D V K N T I N C D F S K Y V P P P P P P K P T P P T P P T P P T P P T P P T P P T P P T P R P V H N R K V M F F V A G A V L V A I L I S T V R W\n")
        f.write("M A S N T V S A Q G G S N R P V R D F S N I Q D V A Q F L L F D P I W N E Q P G S I V P W K M N R E Q A L A E R Y P E L Q T S E P S E D Y S G P V E S L E L L P L E I K L D I M Q Y L S W E Q I S W C K H P W L W T R W Y K D N V V R V S A I T F E D F Q R E Y A F P E K I Q E I H F T D T R A E E I K A I L E T T P N V T R L V I R R I D D M N Y N T H G D L G L D D L E F L T H L M V E D A C G F T D F W A P S L T H L T I K N L D M H P R W F G P V M D G I K S M Q S T L K Y L Y I F E T Y G V N K P F V Q W C T D N I E T F Y C T N S Y R Y E N V P R P I Y V W V L F Q E D E W H G Y R V E D N K F H R R Y M Y S T I L H K R D T D W V E N N P L K T P A Q V E M Y K F L L R I S Q L N R D G T G Y E S D S D P E N E H F D D E S F S S G E E D S S D E D D P T W A P D S D D S D W E T E T E E E P S V A A R I L E K G K L T I T N L M K S L G F K P K P K K I Q S I D R Y F C S L D S N Y N S E D E D F E Y D S D S E D D D S D S E D D C\n")
        f.write("M Y Q A I N P C P Q S W Y G S P Q L E R E I V C K M S G A P H Y P N Y Y P V H P N A L G G A W F D T S L N A R S L T T T P S L T T C T P P S L A A C T P P T S L G M V D S P P H I N P P R R I G T L C F D F G S A K S P Q R C E C V A S D R P S T T S N T A P D T Y R L L I T N S K T R K N N Y G T C R L E P L T Y G I\n")
        f.write("M A R P L L G K T S S V R R R L E S L S A C S I F F F L R K F C Q K M A S L V F L N S P V Y Q M S N I L L T E R R Q V D R A M G G S D D D G V M V V A L S P S D F K T V L G S A L L A V E R D M V H V V P K Y L Q T P G I L H D M L V L L T P I F G E A L S V D M S G A T D V M V Q Q I A T A G F V D V D P L H S S V S W K D N V S C P V A L L A V S N A V R T M M G Q P C Q V T L I I D V G T Q N I L R D L V N L P V E M S G D L Q V M A Y T K D P L G K V P A V G V S V F D S G S V Q K G D A H S V G A P D G L V S F H T H P V S S A V E L N Y H A G W P S N V D M S S L L T M K N L M H V V V A E E G L W T M A R T L S M Q R L T K V L T D A E K D V M R A A A F N L F L P L N E L R V M G T K D S N N K S L K T Y F E V F E T F T I G A L M K H S G V T P T A F V D R R W L D N T I Y H M G F I P W G R D M R F V V E Y D L D G T N P F L N T V P T L M S V K R K A K I Q E M F D N M V S R M V T S\n")
        f.write("M N A K Y D T D Q G V G R M L F L G T I G L A V V V G G L M A Y G Y Y Y D G K T P S S G T S F H T A S P S F S S R Y R Y\n")
        f.write("M R Y T V L I A L Q G A L L L L L L I D D G Q G Q S P Y P Y P G M P C N S S R Q C G L G T C V H S R C A H C S S D G T L C S P E D P T M V W P C C P E S S C Q L V V G L P S L V N H Y N C L P N Q C T D S S Q C P G G F G C M T R R S K C E L C K A D G E A C N S P Y L D W R K D K E C C S G Y C H T E A R G L E G V C I D P K K I F C T P K N P W Q L A P Y P P S Y H Q P T T L R P P T S L Y D S W L M S G F L V K S T T A P S T Q E E E D D Y\n")
        f.write("M Q N P L P E V M S P E H D K R T T T P M S K E A N K F I R E L D K K P G D L A V V S D F V K R N T G K R L P I G K R S N L Y V R I C D L S G T I Y M G E T F I L E S W E E L Y L P E P T K M E V L G T L E S C C G I P P F P E W I V M V G E D Q C V Y A Y G D E E I L L F A Y S V K Q L V E E G I Q E T G I S Y K Y P D D I S D V D E E V L Q Q D E E I Q K I R K K T R E F V D K D A Q E F Q D F L N S L D A S L L S\n")
        f.write("M D S L N E V C Y E Q I K G T F Y K G L F G D F P L I V D K K T G C F N A T K L C V L G G K R F V D W N K T L R S K K L I Q Y Y E T R C D I K T E S L L Y E I K G D N N D E I T K Q I T G T Y L P K E F I L D I A S W I S V E F Y D K C N N I I I N Y F V N E Y K T M D K K T L Q S K I N E V E E K M Q K L L N E K E E E L Q E K N D K I D E L I L F S K R M E E D R K K D R E M M I K Q E K M L R E L G I H L E D V S S Q N N E L I E K V D E Q V E Q N A V L N F K I D N I Q N K L E I A V E D R A P Q P K Q N L K R E R F I L L K R N D D Y Y P Y Y T I R A Q D I N A R S A L K R Q K N L Y N E V S V L L D L T C H P N S K T L Y V R V K D E L K Q K G V V F N L C K V S I S N S K I N E E E L I K A M E T I N D E K R D V\n")

fake_data()

spm.SentencePieceTrainer.train(input='vocab.txt',
                               model_prefix='protein',
                               vocab_size=30,
                               model_type="word",
                               #user_defined_symbols="<MASK>",
                               pad_id=0,
                               unk_id=1,
                               bos_id=2,
                               eos_id=3,
                               pad_piece="[PAD]",
                               unk_piece="[UNK]",
                               bos_piece="[BOS]",
                               eos_piece="[EOS]")

tokenizer = spm.SentencePieceProcessor(model_file='protein.model')

train_files = glob.glob("dataset/train*",recursive=True)
random.shuffle(train_files)

def mask_seq(seq,mask_prob=0.15):
  seq = np.array(seq)
  minValue = 1
  maxValue = len(seq) - 2
  max_mask_tokens = int(maxValue * 0.15 + 0.5)
  randomlist = random.sample(range(minValue, maxValue), max_mask_tokens)
  seq_masked = seq
  seq_masked[randomlist] = tokenizer.encode("[MASK]")[0]
  return seq_masked

def get_seq(train_files):
  while True:
    for file in train_files:
      with open(file) as fp:
        for line in fp:
          yield line

def get_batch(seq_gen, batch_length):

  batch = []
  while True:
    seq = next(seq_gen)
    seq_ids = tokenizer.encode(seq,add_bos=True,add_eos=True)

    new_batch_len = len(batch) + len(seq_ids)

    if new_batch_len <= batch_length :
      batch = batch + seq_ids
      continue

    next_batch = batch
    batch = seq_ids

    yield next_batch

# Set up the data pipeline.
def my_inputs(n_devices):

  MAX_BATCH_LENGTH = 1024*4

  seq_gen = get_seq(train_files)
  batch_gen = get_batch(seq_gen,MAX_BATCH_LENGTH)

  while True:

    inputs = []
    targets = []
    mask = []
    for i in range(n_devices):
      batch_ids = next(batch_gen)
      masked_seq_ids = mask_seq(batch_ids)
      pad_amount = MAX_BATCH_LENGTH - len(batch_ids)

      inputs.append(np.pad(masked_seq_ids, (0,pad_amount)))
      targets.append(np.pad(batch_ids, (0,pad_amount)))
      mask.append(np.pad(np.ones_like(batch_ids, dtype=np.float32),
                          (0,pad_amount),
                          mode='constant'))
    inputs = np.stack(inputs)
    targets = np.stack(targets)
    mask = np.stack(mask)
    yield (inputs, targets, mask)

inp_gen_test = my_inputs(trax.fastmath.device_count())

res = next(inp_gen_test)

print(tokenizer.decode(res[0][0].tolist()))

print(tokenizer.decode(res[1][0].tolist()))

# Configure hyperparameters.
gin.parse_config("""
import trax.layers
import trax.models
import trax.optimizers
import trax.data.inputs
import trax.supervised.trainer_lib

# Parameters that will vary between experiments:
# ==============================================================================
train.model = @trax.models.Reformer
n_layers = 15
n_heads = 16
dropout = 0.1
n_tokens = 40000 # They have used very small n_tokens = 2048
vocab_size= 30
d_model = 1024
d_ff = 4096

# Done
# Parameters for MultifactorSchedule:
# ==============================================================================
multifactor.constant = 0.088
multifactor.decay_factor = 0.5
multifactor.factors = 'constant * linear_warmup * rsqrt_decay'
multifactor.steps_per_cycle = 100000
multifactor.steps_per_decay = 20000
multifactor.warmup_steps = 8000

# Done
# Parameters for Adam:
# ==============================================================================
Adam.b1 = 0.9
Adam.b2 = 0.98
Adam.eps = 1e-09
Adam.weight_decay_rate = 1e-05

# Done
# Parameters for SelfAttention:
# ==============================================================================
#trax.layers.SelfAttention.attention_dropout = 0.05
#trax.layers.SelfAttention.chunk_len = 64
#trax.layers.SelfAttention.n_chunks_before = 1
#trax.layers.SelfAttention.n_parallel_heads = 1
trax.layers.SelfAttention.causal = False
trax.layers.SelfAttention.chunk_len = None
trax.layers.SelfAttention.masked = False
trax.layers.SelfAttention.n_chunks_after = 0
trax.layers.SelfAttention.n_chunks_before = 0
trax.layers.SelfAttention.n_parallel_heads = None
trax.layers.SelfAttention.predict_drop_len = 64
trax.layers.SelfAttention.predict_mem_len = 192
trax.layers.SelfAttention.share_qk = False
trax.layers.SelfAttention.use_python_loop = False
trax.layers.SelfAttention.use_reference_code = False

# Done
# Parameters for EncDecAttention:
# ==============================================================================
trax.layers.EncDecAttention.masked = True
trax.layers.EncDecAttention.n_parallel_heads = None
trax.layers.EncDecAttention.use_python_loop = False
trax.layers.EncDecAttention.use_reference_code = False

# Done
# Parameters for LSHSelfAttention:
# ==============================================================================
#LSHSelfAttention.attention_dropout = 0.0
#LSHSelfAttention.chunk_len = 64
#LSHSelfAttention.n_buckets = [64, 128]
#LSHSelfAttention.n_chunks_after = 0
#LSHSelfAttention.n_chunks_before = 1
#LSHSelfAttention.n_hashes = 1
#LSHSelfAttention.n_parallel_heads = 1
#LSHSelfAttention.predict_drop_len = 128
#LSHSelfAttention.predict_mem_len = 1024

# Done
# Parameters for Reformer:
# ==============================================================================
Reformer.d_model = %d_model
Reformer.d_ff = %d_ff
Reformer.dropout = %dropout
Reformer.ff_activation = @trax.layers.Relu
Reformer.max_len = %n_tokens
Reformer.mode = 'train'
Reformer.n_heads = %n_heads
Reformer.n_encoder_layers = %n_layers
Reformer.n_decoder_layers = %n_layers
Reformer.input_vocab_size = %vocab_size

""")

# Set up a Trainer.
output_dir = os.path.expanduser('train_dir/')

trainer = trax.supervised.Trainer(
    model=trax.models.Reformer,
    loss_fn=trax.layers.CrossEntropyLoss(),
    optimizer=trax.optimizers.Adam,
    lr_schedule=trax.lr.multifactor(),
    inputs=trax.data.inputs.Inputs(my_inputs),
    output_dir=output_dir)

# Run one training step, to make sure the model fits in memory.
# The first time trainer.train_epoch is called, it will JIT the entire network
# architecture, which takes around 2 minutes. The JIT-compiled model is saved
# so subsequent runs will be much faster than the first.
trainer.train_epoch(n_steps=1, n_eval_steps=1)

Any idea how I can solve this issue ?

I have test it with a Google cloud machine with 8x V100, and it gives a similar issue:
2020-08-27 18:34:07.848853: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:55]
********************************
Very slow compile?  If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
********************************

For the GPU problem, I found out that Jax doesn't detect the GPU because of Cuda 11 problem with tensorflow.
Still, I can't figure out what is the problem with the TPU.

I think it may just be taking a very long time to compile. (And potentially the tpu_driver has a time-out that's too short here).

Thanks @jekbradbury Is there a solution around it ?

I have fixed the GPU issue with Tensorflow but Jax still can't see the GPU:
>>> import tensorflow as tf
2020-08-27 20:15:19.541390: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudart.so.10.1
>>> tf.test.gpu_device_name()
2020-08-27 20:15:25.408895: I tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN)to use the following CPU instructions in performance-critical operations:  AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2020-08-27 20:15:25.418379: I tensorflow/core/platform/profile_utils/cpu_utils.cc:104] CPU Frequency: 2300000000 Hz
2020-08-27 20:15:25.419979: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x56287ad69760 initialized for platform Host (this does not guarantee that XLA will be used). Devices:
2020-08-27 20:15:25.420015: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (0): Host, Default Version
2020-08-27 20:15:25.422701: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcuda.so.1
2020-08-27 20:15:30.000531: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:15:30.072938: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:15:30.096370: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:15:30.119383: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:15:30.144339: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:15:30.166636: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:15:30.192891: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:15:30.223290: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:15:30.225406: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x56287add1100 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2020-08-27 20:15:30.225441: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (0): Tesla V100-SXM2-16GB, Compute Capability 7.0
2020-08-27 20:15:30.225448: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (1): Tesla V100-SXM2-16GB, Compute Capability 7.0
2020-08-27 20:15:30.225461: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (2): Tesla V100-SXM2-16GB, Compute Capability 7.0
2020-08-27 20:15:30.225476: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (3): Tesla V100-SXM2-16GB, Compute Capability 7.0
2020-08-27 20:15:30.225491: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (4): Tesla V100-SXM2-16GB, Compute Capability 7.0
2020-08-27 20:15:30.225504: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (5): Tesla V100-SXM2-16GB, Compute Capability 7.0
2020-08-27 20:15:30.225517: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (6): Tesla V100-SXM2-16GB, Compute Capability 7.0
2020-08-27 20:15:30.225530: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (7): Tesla V100-SXM2-16GB, Compute Capability 7.0
2020-08-27 20:15:30.266243: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:15:30.267744: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1716] Found device 0 with properties:
pciBusID: 0000:00:04.0 name: Tesla V100-SXM2-16GB computeCapability: 7.0
coreClock: 1.53GHz coreCount: 80 deviceMemorySize: 15.78GiB deviceMemoryBandwidth: 836.37GiB/s
2020-08-27 20:15:30.267909: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:15:30.269364: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1716] Found device 1 with properties:
pciBusID: 0000:00:05.0 name: Tesla V100-SXM2-16GB computeCapability: 7.0
coreClock: 1.53GHz coreCount: 80 deviceMemorySize: 15.78GiB deviceMemoryBandwidth: 836.37GiB/s
2020-08-27 20:15:30.269468: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:15:30.271025: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1716] Found device 2 with properties:
pciBusID: 0000:00:06.0 name: Tesla V100-SXM2-16GB computeCapability: 7.0
coreClock: 1.53GHz coreCount: 80 deviceMemorySize: 15.78GiB deviceMemoryBandwidth: 836.37GiB/s
2020-08-27 20:15:30.271152: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:15:30.272581: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1716] Found device 3 with properties:
pciBusID: 0000:00:07.0 name: Tesla V100-SXM2-16GB computeCapability: 7.0
coreClock: 1.53GHz coreCount: 80 deviceMemorySize: 15.78GiB deviceMemoryBandwidth: 836.37GiB/s
2020-08-27 20:15:30.272679: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:15:30.274168: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1716] Found device 4 with properties:
pciBusID: 0000:00:08.0 name: Tesla V100-SXM2-16GB computeCapability: 7.0
coreClock: 1.53GHz coreCount: 80 deviceMemorySize: 15.78GiB deviceMemoryBandwidth: 836.37GiB/s
2020-08-27 20:15:30.274298: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:15:30.275798: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1716] Found device 5 with properties:
pciBusID: 0000:00:09.0 name: Tesla V100-SXM2-16GB computeCapability: 7.0
coreClock: 1.53GHz coreCount: 80 deviceMemorySize: 15.78GiB deviceMemoryBandwidth: 836.37GiB/s
2020-08-27 20:15:30.275909: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:15:30.277374: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1716] Found device 6 with properties:
pciBusID: 0000:00:0a.0 name: Tesla V100-SXM2-16GB computeCapability: 7.0
coreClock: 1.53GHz coreCount: 80 deviceMemorySize: 15.78GiB deviceMemoryBandwidth: 836.37GiB/s
2020-08-27 20:15:30.277469: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:15:30.278953: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1716] Found device 7 with properties:
pciBusID: 0000:00:0b.0 name: Tesla V100-SXM2-16GB computeCapability: 7.0
coreClock: 1.53GHz coreCount: 80 deviceMemorySize: 15.78GiB deviceMemoryBandwidth: 836.37GiB/s
2020-08-27 20:15:30.279005: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudart.so.10.1
2020-08-27 20:15:30.280902: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcublas.so.10
2020-08-27 20:15:30.282711: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcufft.so.10
2020-08-27 20:15:30.283056: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcurand.so.10
2020-08-27 20:15:30.285176: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcusolver.so.10
2020-08-27 20:15:30.286358: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcusparse.so.10
2020-08-27 20:15:30.286555: W tensorflow/stream_executor/platform/default/dso_loader.cc:59] Could not load dynamic library 'libcudnn.so.7'; dlerror: libcudnn.so.7: cannot open shared object file: No such file or directory
2020-08-27 20:15:30.286576: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1753] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...
2020-08-27 20:15:30.286671: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1257] Device interconnect StreamExecutor with strength 1 edge matrix:
2020-08-27 20:15:30.286687: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1263]      0 1 2 3 4 5 6 7
2020-08-27 20:15:30.286697: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1276] 0:   N Y Y Y N Y N N
2020-08-27 20:15:30.286708: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1276] 1:   Y N Y Y N N N N
2020-08-27 20:15:30.286718: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1276] 2:   Y Y N Y N N N Y
2020-08-27 20:15:30.286733: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1276] 3:   Y Y Y N N N N N
2020-08-27 20:15:30.286746: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1276] 4:   N N N N N Y Y Y
2020-08-27 20:15:30.286756: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1276] 5:   Y N N N Y N Y Y
2020-08-27 20:15:30.286767: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1276] 6:   N N N N Y Y N Y
2020-08-27 20:15:30.286775: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1276] 7:   N N Y N Y Y Y N
''
>>> from jax.lib import xla_bridge
>>> print(xla_bridge.get_backend().platform)
/home/ahmed/anaconda3/envs/trax/lib/python3.6/site-packages/jax/lib/xla_bridge.py:130: UserWarning: No GPU/TPU found, falling back to CPU.
  warnings.warn('No GPU/TPU found, falling back to CPU.')
cpu
>>>

I have fixed the Tensorflow issue but still Jax can't see the GPUs:
Python 3.6.10 |Anaconda, Inc.| (default, May  8 2020, 02:54:21)
[GCC 7.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import tensorflow as tf
2020-08-27 20:21:00.806963: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudart.so.10.1
>>> physical_devices = tf.config.list_physical_devices('GPU')
2020-08-27 20:21:04.095189: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcuda.so.1
2020-08-27 20:21:04.548746: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:21:04.550269: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1716] Found device 0 with properties:
pciBusID: 0000:00:04.0 name: Tesla V100-SXM2-16GB computeCapability: 7.0
coreClock: 1.53GHz coreCount: 80 deviceMemorySize: 15.78GiB deviceMemoryBandwidth: 836.37GiB/s
2020-08-27 20:21:04.550419: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:21:04.551886: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1716] Found device 1 with properties:
pciBusID: 0000:00:05.0 name: Tesla V100-SXM2-16GB computeCapability: 7.0
coreClock: 1.53GHz coreCount: 80 deviceMemorySize: 15.78GiB deviceMemoryBandwidth: 836.37GiB/s
2020-08-27 20:21:04.551995: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:21:04.553494: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1716] Found device 2 with properties:
pciBusID: 0000:00:06.0 name: Tesla V100-SXM2-16GB computeCapability: 7.0
coreClock: 1.53GHz coreCount: 80 deviceMemorySize: 15.78GiB deviceMemoryBandwidth: 836.37GiB/s
2020-08-27 20:21:04.553588: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:21:04.555113: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1716] Found device 3 with properties:
pciBusID: 0000:00:07.0 name: Tesla V100-SXM2-16GB computeCapability: 7.0
coreClock: 1.53GHz coreCount: 80 deviceMemorySize: 15.78GiB deviceMemoryBandwidth: 836.37GiB/s
2020-08-27 20:21:04.555202: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:21:04.556676: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1716] Found device 4 with properties:
pciBusID: 0000:00:08.0 name: Tesla V100-SXM2-16GB computeCapability: 7.0
coreClock: 1.53GHz coreCount: 80 deviceMemorySize: 15.78GiB deviceMemoryBandwidth: 836.37GiB/s
2020-08-27 20:21:04.556769: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:21:04.558275: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1716] Found device 5 with properties:
pciBusID: 0000:00:09.0 name: Tesla V100-SXM2-16GB computeCapability: 7.0
coreClock: 1.53GHz coreCount: 80 deviceMemorySize: 15.78GiB deviceMemoryBandwidth: 836.37GiB/s
2020-08-27 20:21:04.558362: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:21:04.559834: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1716] Found device 6 with properties:
pciBusID: 0000:00:0a.0 name: Tesla V100-SXM2-16GB computeCapability: 7.0
coreClock: 1.53GHz coreCount: 80 deviceMemorySize: 15.78GiB deviceMemoryBandwidth: 836.37GiB/s
2020-08-27 20:21:04.559919: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:21:04.561362: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1716] Found device 7 with properties:
pciBusID: 0000:00:0b.0 name: Tesla V100-SXM2-16GB computeCapability: 7.0
coreClock: 1.53GHz coreCount: 80 deviceMemorySize: 15.78GiB deviceMemoryBandwidth: 836.37GiB/s
2020-08-27 20:21:04.561405: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudart.so.10.1
2020-08-27 20:21:04.563455: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcublas.so.10
2020-08-27 20:21:04.565374: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcufft.so.10
2020-08-27 20:21:04.565677: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcurand.so.10
2020-08-27 20:21:04.567552: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcusolver.so.10
2020-08-27 20:21:04.568606: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcusparse.so.10
2020-08-27 20:21:04.572985: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudnn.so.7
2020-08-27 20:21:04.573151: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:21:04.574731: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:21:04.576168: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:21:04.577606: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:21:04.579063: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:21:04.580476: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:21:04.582068: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:21:04.583612: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:21:04.585144: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:21:04.586596: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:21:04.588028: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:21:04.589484: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:21:04.590972: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:21:04.592478: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:21:04.593962: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:21:04.595447: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2020-08-27 20:21:04.596795: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1858] Adding visible gpu devices: 0, 1, 2, 3, 4, 5, 6, 7
>>> from jax.lib import xla_bridge
>>> print(xla_bridge.get_backend().platform)
/home/ahmed/anaconda3/envs/trax/lib/python3.6/site-packages/jax/lib/xla_bridge.py:130: UserWarning: No GPU/TPU found, falling back to CPU.
  warnings.warn('No GPU/TPU found, falling back to CPU.')
cpu
>>>

Any news about this issue ?

If you just import jax and not TF, does jax see the GPUs?

No, it doesn't see the GPU, but TensorFlow sees it.

Can you tell which line of code is causing the long compile?
For the GPU issue, how are you installing jax+ jaxib?

Can you tell which line of code is causing the long compile?
The trainer:

trainer = trax.supervised.Trainer(
    model=trax.models.Reformer,
    loss_fn=trax.layers.CrossEntropyLoss(),
    optimizer=trax.optimizers.Adam,
    lr_schedule=trax.lr.multifactor(),
    inputs=trax.data.inputs.Inputs(my_inputs),
    output_dir=output_dir)

I believe this trigger JAX compilation.

For the GPU issue, I figured out the problem. I was using :
pip install --upgrade jax jaxlib
Which doesn't support CUDA, and I followed the readme and Jax now can see the GPUs.
I think the default Jax installation should support both TPUs and GPUs. This will be easier for users.
The issue of very long compilation still exists using the GPU for the same command mentioned above.

During this step, I can see that only a single core is being used for the trainer function:

I have tested the above script using GPUs and it is working without any issue. The long compilation process only occurs with TPUs.

Glad to hear you resolved the jaxlib issue! Unfortunately we cannot support GPUs by default since we need different jaxlibs for different CUDA versions. https://github.com/google/jax/pull/4065 should make this a little simpler at least.
For the long compilation on TPU, like @jekbradbury says, there might be an issue with the underlying TPU driver timing out. We're aware of this issue, but I don't know of any immediate workaround. Your best bet is probably to use GPU for now.

Thanks a lot @skye  for your explanation and support.
I will use the GPUs for now and I hope the TPU issue will be solved in the near future.

I am also getting the same error, where I am trying to train the standard TrnasformerLM, with the slight modification in the parameters(vocab_size=50000, max_len=1024). The same code works on the COLAB + TPU setup, but when I try a bigger TPU with GCE-VM + CLOUD-TPU setup I am getting the same error.

now 12.01.22 and the problem still exists

I am also getting the same error, where I am trying to train the standard TrnasformerLM, with the slight modification in the parameters(vocab_size=50000, max_len=1024). The same code works on the COLAB + TPU setup, but when I try it with the GCE-VM + CLOUD-TPU setup I am getting the same error.

I have the same issue. I can run big bird model with max_length of 1024 but compilation times out with max length of 2048.

authors: ProBuro agemagician jayendra13 jekbradbury sarthak-314 skye