Falling back to the CUDA driver for PTX compilation; ptxas does not support CC 8.9

Problems between tensorflow, cuda, jaxlib... I don't know how to solve... I'm using CUDA 12.3.0 and nvidia/cuda:12.3-base-ubuntu20.04

```
I0626 18:28:37.506076 138088085884928 run_docker.py:116] Mounting /mnt/sdb1/backup/alphafold_databases/uniref90 -> /mnt/sdb1/backup/alphafold_mountpoint/uniref90_database_path
I0626 18:28:37.506205 138088085884928 run_docker.py:116] Mounting /mnt/sdb1/backup/alphafold_databases/mgnify -> /mnt/sdb1/backup/alphafold_mountpoint/mgnify_database_path
I0626 18:28:37.506297 138088085884928 run_docker.py:116] Mounting /mnt/sdb1/backup/alphafold_databases -> /mnt/sdb1/backup/alphafold_mountpoint/data_dir
I0626 18:28:37.506383 138088085884928 run_docker.py:116] Mounting /mnt/sdb1/backup/alphafold_databases/pdb_mmcif/mmcif_files -> /mnt/sdb1/backup/alphafold_mountpoint/template_mmcif_dir
I0626 18:28:37.506479 138088085884928 run_docker.py:116] Mounting /mnt/sdb1/backup/alphafold_databases/pdb_mmcif -> /mnt/sdb1/backup/alphafold_mountpoint/obsolete_pdbs_path
I0626 18:28:37.506588 138088085884928 run_docker.py:116] Mounting /mnt/sdb1/backup/alphafold_databases/pdb70 -> /mnt/sdb1/backup/alphafold_mountpoint/pdb70_database_path
I0626 18:28:37.506686 138088085884928 run_docker.py:116] Mounting /mnt/sdb1/backup/alphafold_databases/small_bfd -> /mnt/sdb1/backup/alphafold_mountpoint/small_bfd_database_path
I0626 18:28:42.291707 138088085884928 run_docker.py:258] I0626 21:28:42.290248 133795469678400 templates.py:858] Using precomputed obsolete pdbs /mnt/sdb1/backup/alphafold_mountpoint/obsolete_pdbs_path/obsolete.dat.
I0626 18:28:44.780305 138088085884928 run_docker.py:258] I0626 21:28:44.779594 133795469678400 xla_bridge.py:353] Unable to initialize backend 'tpu_driver': NOT_FOUND: Unable to find driver in registry given worker:
I0626 18:28:44.968280 138088085884928 run_docker.py:258] I0626 21:28:44.967394 133795469678400 xla_bridge.py:353] Unable to initialize backend 'rocm': NOT_FOUND: Could not find registered platform with name: "rocm". Available platform names are: Interpreter Host CUDA
I0626 18:28:44.968665 138088085884928 run_docker.py:258] I0626 21:28:44.968128 133795469678400 xla_bridge.py:353] Unable to initialize backend 'tpu': module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'
I0626 18:28:44.968835 138088085884928 run_docker.py:258] I0626 21:28:44.968231 133795469678400 xla_bridge.py:353] Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.
I0626 18:28:58.638173 138088085884928 run_docker.py:258] I0626 21:28:58.637206 133795469678400 run_alphafold.py:524] Have 5 models: ['model_1_pred_0', 'model_2_pred_0', 'model_3_pred_0', 'model_4_pred_0', 'model_5_pred_0']
I0626 18:28:58.638509 138088085884928 run_docker.py:258] I0626 21:28:58.637522 133795469678400 run_alphafold.py:538] Using random seed 703549940514398884 for the data pipeline
I0626 18:28:58.638677 138088085884928 run_docker.py:258] I0626 21:28:58.638001 133795469678400 run_alphafold.py:245] Predicting teste
I0626 18:28:58.639462 138088085884928 run_docker.py:258] I0626 21:28:58.638904 133795469678400 jackhmmer.py:133] Launching subprocess "/usr/bin/jackhmmer -o /dev/null -A /tmp/tmpuc2_2bha/output.sto --noali --F1 0.0005 --F2 5e-05 --F3 5e-07 --incE 0.0001 -E 0.0001 --cpu 8 -N 1 /mnt/sdb1/backup/alphafold_mountpoint/fasta_path_0/teste.faa /mnt/sdb1/backup/alphafold_mountpoint/uniref90_database_path/uniref90.fasta"
I0626 18:28:58.640388 138088085884928 run_docker.py:258] I0626 21:28:58.639911 133795469678400 utils.py:36] Started Jackhmmer (uniref90.fasta) query
I0626 18:39:55.007368 138088085884928 run_docker.py:258] I0626 21:39:55.006375 133795469678400 utils.py:40] Finished Jackhmmer (uniref90.fasta) query in 656.366 seconds
I0626 18:40:01.191493 138088085884928 run_docker.py:258] I0626 21:40:01.190375 133795469678400 jackhmmer.py:133] Launching subprocess "/usr/bin/jackhmmer -o /dev/null -A /tmp/tmpllsuut3p/output.sto --noali --F1 0.0005 --F2 5e-05 --F3 5e-07 --incE 0.0001 -E 0.0001 --cpu 8 -N 1 /mnt/sdb1/backup/alphafold_mountpoint/fasta_path_0/teste.faa /mnt/sdb1/backup/alphafold_mountpoint/mgnify_database_path/mgy_clusters_2022_05.fa"
I0626 18:40:01.192113 138088085884928 run_docker.py:258] I0626 21:40:01.191634 133795469678400 utils.py:36] Started Jackhmmer (mgy_clusters_2022_05.fa) query
I0626 18:56:02.226763 138088085884928 run_docker.py:258] I0626 21:56:02.225653 133795469678400 utils.py:40] Finished Jackhmmer (mgy_clusters_2022_05.fa) query in 961.034 seconds
I0626 18:56:10.551476 138088085884928 run_docker.py:258] I0626 21:56:10.550518 133795469678400 hhsearch.py:85] Launching subprocess "/usr/bin/hhsearch -i /tmp/tmpcedwwkll/query.a3m -o /tmp/tmpcedwwkll/output.hhr -maxseq 1000000 -d /mnt/sdb1/backup/alphafold_mountpoint/pdb70_database_path/pdb70"
I0626 18:56:10.552288 138088085884928 run_docker.py:258] I0626 21:56:10.551964 133795469678400 utils.py:36] Started HHsearch query
I0626 18:57:57.670601 138088085884928 run_docker.py:258] I0626 21:57:57.669543 133795469678400 utils.py:40] Finished HHsearch query in 107.117 seconds
I0626 18:58:02.625957 138088085884928 run_docker.py:258] I0626 21:58:02.625021 133795469678400 jackhmmer.py:133] Launching subprocess "/usr/bin/jackhmmer -o /dev/null -A /tmp/tmpt6dq4gsr/output.sto --noali --F1 0.0005 --F2 5e-05 --F3 5e-07 --incE 0.0001 -E 0.0001 --cpu 8 -N 1 /mnt/sdb1/backup/alphafold_mountpoint/fasta_path_0/teste.faa /mnt/sdb1/backup/alphafold_mountpoint/small_bfd_database_path/bfd-first_non_consensus_sequences.fasta"
I0626 18:58:02.626496 138088085884928 run_docker.py:258] I0626 21:58:02.626078 133795469678400 utils.py:36] Started Jackhmmer (bfd-first_non_consensus_sequences.fasta) query
I0626 19:00:15.475680 138088085884928 run_docker.py:258] I0626 22:00:15.474710 133795469678400 utils.py:40] Finished Jackhmmer (bfd-first_non_consensus_sequences.fasta) query in 132.849 seconds
I0626 19:00:17.375041 138088085884928 run_docker.py:258] I0626 22:00:17.373987 133795469678400 templates.py:879] Searching for template for: MDGQNPKPLPLTSLLLPVSIAVDLVARKLVWADATRGTIESLDLVTVFTGTPFIVQQIKSHVSSVSAAHNEIHWTSRNNSSLECIDISAEPHVRRSVLLPTGRNGSYSRRVLIASQVPDFEPGPCGRNNGGCSHTCLPVRTTERACFCPPGMALGTNNITCRVENGTCRPHELPCAGGCIAATYWCDGHKDCADNADEATCGTTCPPKDFTCQNGKCIDTAWRCDGYDDCDDQSDEANCPYRTCASDQFNCKSGACLPFYWRCDGANDCPDGDDELDCRTLRCPNGHDRCANGQCIPRDWSCDGHADCSDSSDETNCTETVSCFEDDFHCANGQCIDKRLRCDHDEDCEDNSDESGCDYAPANNSKCVKGMVGCGDGQCVYTHDMCDGYADCHNGWDEHNCSAPVCQSAEFFCPGTKRCILQSWLCDGDDDCGDAMDELLARCRPTTLPPPTDAPCWSDQFQCGSHECIAWSSVCDGRSDCADFSDEGSHCEHHCATANGGCAHICRESPSGPQCSCRPGYRLNNDRKSCDDIDECLSPGHCSHFCQNSKGGYK
I0626 19:00:17.948528 138088085884928 run_docker.py:258] I0626 22:00:17.947738 133795469678400 templates.py:267] Found an exact template match 3m0c_C.
I0626 19:00:18.206189 138088085884928 run_docker.py:258] I0626 22:00:18.205397 133795469678400 templates.py:267] Found an exact template match 1n7d_A.
I0626 19:00:18.226388 138088085884928 run_docker.py:258] I0626 22:00:18.225925 133795469678400 templates.py:267] Found an exact template match 3m0c_C.
I0626 19:00:18.244172 138088085884928 run_docker.py:258] I0626 22:00:18.243780 133795469678400 templates.py:267] Found an exact template match 1n7d_A.
I0626 19:00:18.264382 138088085884928 run_docker.py:258] I0626 22:00:18.264040 133795469678400 templates.py:267] Found an exact template match 3m0c_C.
I0626 19:00:18.280781 138088085884928 run_docker.py:258] I0626 22:00:18.280326 133795469678400 templates.py:267] Found an exact template match 3m0c_C.
I0626 19:00:18.296232 138088085884928 run_docker.py:258] I0626 22:00:18.295779 133795469678400 templates.py:913] Skipped invalid hit 3M0C_C Proprotein convertase subtilisin/kexin type 9; PROTEIN COMPLEX, BETA PROPELLER, CHOLESTEROL; 7.01A {Homo sapiens}, error: None, warning: 3m0c_C (sum_probs: 150.4, rank: 5): feature extracting errors: Template all atom mask was all zeros: 3m0c_C. Residue range: 21-225, mmCIF parsing errors: {}
I0626 19:00:18.296375 138088085884928 run_docker.py:258] I0626 22:00:18.296076 133795469678400 templates.py:267] Found an exact template match 1n7d_A.
I0626 19:00:18.915404 138088085884928 run_docker.py:258] I0626 22:00:18.914566 133795469678400 templates.py:267] Found an exact template match 4a0p_A.
I0626 19:00:19.795932 138088085884928 run_docker.py:258] I0626 22:00:19.795142 133795469678400 templates.py:267] Found an exact template match 6h15_B.
I0626 19:00:20.890163 138088085884928 run_docker.py:258] I0626 22:00:20.889332 133795469678400 templates.py:267] Found an exact template match 5b4x_D.
I0626 19:00:21.271526 138088085884928 run_docker.py:258] I0626 22:00:21.270835 133795469678400 templates.py:267] Found an exact template match 3p5c_L.
I0626 19:00:21.617020 138088085884928 run_docker.py:258] I0626 22:00:21.616351 133795469678400 templates.py:267] Found an exact template match 3p5b_L.
I0626 19:00:21.630503 138088085884928 run_docker.py:258] I0626 22:00:21.630137 133795469678400 templates.py:267] Found an exact template match 4a0p_A.
I0626 19:00:21.653079 138088085884928 run_docker.py:258] I0626 22:00:21.652633 133795469678400 templates.py:267] Found an exact template match 6h15_B.
I0626 19:00:21.919118 138088085884928 run_docker.py:258] I0626 22:00:21.918352 133795469678400 templates.py:267] Found an exact template match 4dg6_A.
I0626 19:00:25.451145 138088085884928 run_docker.py:258] I0626 22:00:25.450296 133795469678400 templates.py:267] Found an exact template match 5o32_H.
I0626 19:00:25.463952 138088085884928 run_docker.py:258] I0626 22:00:25.463468 133795469678400 templates.py:267] Found an exact template match 4dg6_A.
I0626 19:00:26.142724 138088085884928 run_docker.py:258] I0626 22:00:26.141979 133795469678400 templates.py:267] Found an exact template match 3s94_A.
I0626 19:00:26.607441 138088085884928 run_docker.py:258] I0626 22:00:26.606480 133795469678400 templates.py:267] Found an exact template match 3sob_B.
I0626 19:00:26.864853 138088085884928 run_docker.py:258] I0626 22:00:26.864096 133795469678400 templates.py:267] Found an exact template match 3sov_A.
I0626 19:00:26.876659 138088085884928 run_docker.py:258] I0626 22:00:26.876217 133795469678400 templates.py:267] Found an exact template match 3s94_A.
I0626 19:00:28.133055 138088085884928 run_docker.py:258] I0626 22:00:28.132189 133795469678400 pipeline.py:234] Uniref90 MSA size: 10000 sequences.
I0626 19:00:28.133450 138088085884928 run_docker.py:258] I0626 22:00:28.132347 133795469678400 pipeline.py:235] BFD MSA size: 5659 sequences.
I0626 19:00:28.133652 138088085884928 run_docker.py:258] I0626 22:00:28.132400 133795469678400 pipeline.py:236] MGnify MSA size: 501 sequences.
I0626 19:00:28.133876 138088085884928 run_docker.py:258] I0626 22:00:28.132448 133795469678400 pipeline.py:237] Final (deduplicated) MSA size: 13266 sequences.
I0626 19:00:28.134085 138088085884928 run_docker.py:258] I0626 22:00:28.132694 133795469678400 pipeline.py:239] Total number of templates (NB: this can include bad templates and is later filtered to top 4): 20.
I0626 19:00:28.249109 138088085884928 run_docker.py:258] I0626 22:00:28.248342 133795469678400 run_alphafold.py:276] Running model model_1_pred_0 on teste
I0626 19:00:33.687656 138088085884928 run_docker.py:258] 2024-06-26 22:00:33.683810: W tensorflow/tsl/framework/cpu_allocator_impl.cc:83] Allocation of 646744032 exceeds 10% of free system memory.
I0626 19:00:33.700572 138088085884928 run_docker.py:258] 2024-06-26 22:00:33.699746: W tensorflow/tsl/framework/cpu_allocator_impl.cc:83] Allocation of 650249744 exceeds 10% of free system memory.
I0626 19:00:34.524705 138088085884928 run_docker.py:258] 2024-06-26 22:00:34.522261: W tensorflow/tsl/framework/cpu_allocator_impl.cc:83] Allocation of 650249744 exceeds 10% of free system memory.
I0626 19:00:34.999829 138088085884928 run_docker.py:258] 2024-06-26 22:00:34.996149: W tensorflow/tsl/framework/cpu_allocator_impl.cc:83] Allocation of 650249744 exceeds 10% of free system memory.
I0626 19:00:35.482020 138088085884928 run_docker.py:258] 2024-06-26 22:00:35.480341: W tensorflow/tsl/framework/cpu_allocator_impl.cc:83] Allocation of 650249744 exceeds 10% of free system memory.
I0626 19:00:35.981338 138088085884928 run_docker.py:258] I0626 22:00:35.980508 133795469678400 model.py:165] Running predict with shape(feat) = {'aatype': (4, 554), 'residue_index': (4, 554), 'seq_length': (4,), 'template_aatype': (4, 4, 554), 'template_all_atom_masks': (4, 4, 554, 37), 'template_all_atom_positions': (4, 4, 554, 37, 3), 'template_sum_probs': (4, 4, 1), 'is_distillation': (4,), 'seq_mask': (4, 554), 'msa_mask': (4, 508, 554), 'msa_row_mask': (4, 508), 'random_crop_to_size_seed': (4, 2), 'template_mask': (4, 4), 'template_pseudo_beta': (4, 4, 554, 3), 'template_pseudo_beta_mask': (4, 4, 554), 'atom14_atom_exists': (4, 554, 14), 'residx_atom14_to_atom37': (4, 554, 14), 'residx_atom37_to_atom14': (4, 554, 37), 'atom37_atom_exists': (4, 554, 37), 'extra_msa': (4, 5120, 554), 'extra_msa_mask': (4, 5120, 554), 'extra_msa_row_mask': (4, 5120), 'bert_mask': (4, 508, 554), 'true_msa': (4, 508, 554), 'extra_has_deletion': (4, 5120, 554), 'extra_deletion_value': (4, 5120, 554), 'msa_feat': (4, 508, 554, 49), 'target_feat': (4, 554, 22)}
I0626 19:00:36.020103 138088085884928 run_docker.py:258] 2024-06-26 22:00:36.019399: W external/org_tensorflow/tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:231] Falling back to the CUDA driver for PTX compilation; ptxas does not support CC 8.9
I0626 19:00:36.020381 138088085884928 run_docker.py:258] 2024-06-26 22:00:36.019477: W external/org_tensorflow/tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:234] Used ptxas at ptxas
I0626 19:00:36.042867 138088085884928 run_docker.py:258] 2024-06-26 22:00:36.042181: E external/org_tensorflow/tensorflow/compiler/xla/stream_executor/cuda/cuda_driver.cc:628] failed to get PTX kernel "shift_right_logical" from module: CUDA_ERROR_NOT_FOUND: named symbol not found
I0626 19:00:36.042988 138088085884928 run_docker.py:258] 2024-06-26 22:00:36.042287: E external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:2153] Execution of replica 0 failed: INTERNAL: Could not find the corresponding function
I0626 19:00:36.043369 138088085884928 run_docker.py:258] Traceback (most recent call last):
I0626 19:00:36.043655 138088085884928 run_docker.py:258] File "/app/alphafold/run_alphafold.py", line 570, in
I0626 19:00:36.043834 138088085884928 run_docker.py:258] app.run(main)
I0626 19:00:36.043925 138088085884928 run_docker.py:258] File "/opt/conda/lib/python3.10/site-packages/absl/app.py", line 312, in run
I0626 19:00:36.044022 138088085884928 run_docker.py:258] _run_main(main, args)
I0626 19:00:36.044106 138088085884928 run_docker.py:258] File "/opt/conda/lib/python3.10/site-packages/absl/app.py", line 258, in _run_main
I0626 19:00:36.044203 138088085884928 run_docker.py:258] sys.exit(main(argv))
I0626 19:00:36.044284 138088085884928 run_docker.py:258] File "/app/alphafold/run_alphafold.py", line 543, in main
I0626 19:00:36.044366 138088085884928 run_docker.py:258] predict_structure(
I0626 19:00:36.044445 138088085884928 run_docker.py:258] File "/app/alphafold/run_alphafold.py", line 284, in predict_structure
I0626 19:00:36.044534 138088085884928 run_docker.py:258] prediction_result = model_runner.predict(processed_feature_dict,
I0626 19:00:36.044613 138088085884928 run_docker.py:258] File "/app/alphafold/alphafold/model/model.py", line 167, in predict
I0626 19:00:36.044701 138088085884928 run_docker.py:258] result = self.apply(self.params, jax.random.PRNGKey(random_seed), feat)
I0626 19:00:36.044782 138088085884928 run_docker.py:258] File "/opt/conda/lib/python3.10/site-packages/jax/_src/random.py", line 132, in PRNGKey
I0626 19:00:36.044874 138088085884928 run_docker.py:258] key = prng.seed_with_impl(impl, seed)
I0626 19:00:36.044963 138088085884928 run_docker.py:258] File "/opt/conda/lib/python3.10/site-packages/jax/_src/prng.py", line 267, in seed_with_impl
I0626 19:00:36.045042 138088085884928 run_docker.py:258] return random_seed(seed, impl=impl)
I0626 19:00:36.045121 138088085884928 run_docker.py:258] File "/opt/conda/lib/python3.10/site-packages/jax/_src/prng.py", line 580, in random_seed
I0626 19:00:36.045211 138088085884928 run_docker.py:258] return random_seed_p.bind(seeds_arr, impl=impl)
I0626 19:00:36.045290 138088085884928 run_docker.py:258] File "/opt/conda/lib/python3.10/site-packages/jax/core.py", line 329, in bind
I0626 19:00:36.045368 138088085884928 run_docker.py:258] return self.bind_with_trace(find_top_trace(args), args, params)
I0626 19:00:36.045447 138088085884928 run_docker.py:258] File "/opt/conda/lib/python3.10/site-packages/jax/core.py", line 332, in bind_with_trace
I0626 19:00:36.045540 138088085884928 run_docker.py:258] out = trace.process_primitive(self, map(trace.full_raise, args), params)
I0626 19:00:36.045619 138088085884928 run_docker.py:258] File "/opt/conda/lib/python3.10/site-packages/jax/core.py", line 712, in process_primitive
I0626 19:00:36.045696 138088085884928 run_docker.py:258] return primitive.impl(*tracers, **params)
I0626 19:00:36.045773 138088085884928 run_docker.py:258] File "/opt/conda/lib/python3.10/site-packages/jax/_src/prng.py", line 592, in random_seed_impl
I0626 19:00:36.045865 138088085884928 run_docker.py:258] base_arr = random_seed_impl_base(seeds, impl=impl)
I0626 19:00:36.045945 138088085884928 run_docker.py:258] File "/opt/conda/lib/python3.10/site-packages/jax/_src/prng.py", line 597, in random_seed_impl_base
I0626 19:00:36.046033 138088085884928 run_docker.py:258] return seed(seeds)
I0626 19:00:36.046113 138088085884928 run_docker.py:258] File "/opt/conda/lib/python3.10/site-packages/jax/_src/prng.py", line 832, in threefry_seed
I0626 19:00:36.046206 138088085884928 run_docker.py:258] lax.shift_right_logical(seed, lax_internal._const(seed, 32)))
I0626 19:00:36.046286 138088085884928 run_docker.py:258] File "/opt/conda/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 515, in shift_right_logical
I0626 19:00:36.046365 138088085884928 run_docker.py:258] return shift_right_logical_p.bind(x, y)
I0626 19:00:36.046445 138088085884928 run_docker.py:258] File "/opt/conda/lib/python3.10/site-packages/jax/core.py", line 329, in bind
I0626 19:00:36.046524 138088085884928 run_docker.py:258] return self.bind_with_trace(find_top_trace(args), args, params)
I0626 19:00:36.046602 138088085884928 run_docker.py:258] File "/opt/conda/lib/python3.10/site-packages/jax/core.py", line 332, in bind_with_trace
I0626 19:00:36.046698 138088085884928 run_docker.py:258] out = trace.process_primitive(self, map(trace.full_raise, args), params)
I0626 19:00:36.046777 138088085884928 run_docker.py:258] File "/opt/conda/lib/python3.10/site-packages/jax/core.py", line 712, in process_primitive
I0626 19:00:36.046856 138088085884928 run_docker.py:258] return primitive.impl(*tracers, **params)
I0626 19:00:36.046934 138088085884928 run_docker.py:258] File "/opt/conda/lib/python3.10/site-packages/jax/_src/dispatch.py", line 115, in apply_primitive
I0626 19:00:36.047012 138088085884928 run_docker.py:258] return compiled_fun(*args)
I0626 19:00:36.047091 138088085884928 run_docker.py:258] File "/opt/conda/lib/python3.10/site-packages/jax/_src/dispatch.py", line 200, in
I0626 19:00:36.047170 138088085884928 run_docker.py:258] return lambda *args, **kw: compiled(*args, **kw)[0]
I0626 19:00:36.047248 138088085884928 run_docker.py:258] File "/opt/conda/lib/python3.10/site-packages/jax/_src/dispatch.py", line 895, in _execute_compiled
I0626 19:00:36.047358 138088085884928 run_docker.py:258] out_flat = compiled.execute(in_flat)
I0626 19:00:36.047441 138088085884928 run_docker.py:258] jaxlib.xla_extension.XlaRuntimeError: INTERNAL: Could not find the corresponding function
```