jax

2024/02/13/notes.org

@@ -15,57 +15,8 @@ generate(serverendpoint,bindings,model,params,text)->text
 
 
 
-* dqn agent
+* jax on tpu 3-8
 
-
-
-
-equinox 0.11.3 requires jax>=0.4.13, but you have jax 0.3.25 which is incompatible.                                 
-
-levanter 1.1 requires tensorstore==0.1.53, but you have tensorstore 0.1.45 which is incompatible.                   
-
-orbax-checkpoint 0.4.4 requires jax>=0.4.9, but you have jax 0.3.25 which is incompatible.                          
-
-chex 0.1.85 requires jax>=0.4.16, but you have jax 0.3.25 which is incompatible.                                    
-
-Successfully installed jax-0.3.25
-Jax 0.3.25 (Nov 15, 2022)
-!pip install flax==0.6.4 jax==0.3.25 jaxlib==0.3.25`
-
-
-
-triton==2.2.0                                                                                                       
-Twisted==22.1.0                                                                                                     
-typeguard==2.13.3                                                                                                   
-typepy==1.3.2                                                                                                       
-typer==0.9.0                                                                                                        
-types-python-dateutil==2.8.19.20240106                                                                              
-typing==3.7.4.3                                                                                                     
-typing-inspect==0.9.0                                                                                               
-typing_extensions==4.9.0                                                                                            
-tzdata==2023.4                                                                                                      
-ubuntu-advantage-tools==8001                                                                                        
-ufw==0.36.1                                                                                                         
-unattended-upgrades==0.1                                                                                            
-uri-template==1.3.0                                                                                                 
-uritemplate==3.0.1                                                                                                  
-urllib3==2.2.0                                                                                                      
-uvicorn==0.23.2                                                                                                     
-virtualenv==20.23.0                                                                                                 
-wadllib==1.3.6                                                                                                      
-wandb==0.16.3                                                                                                       
-wcwidth==0.2.12                                                                                                     
-webcolors==1.13                                                                                                     
-webencodings==0.5.1                                                                                                 
-websocket-client==1.6.4                                                                                             
-websockets==11.0.3                                                                                                  
-Werkzeug==3.0.1                                                                                                     
-widgetsnbextension==4.0.10                                                                                          
-wrapt==1.14.1                                                                                                       
-xmltodict==0.13.0                                                                                                   
-xxhash==3.4.1                                                                                                       
-yapf==0.32.0                                                                                                        
-yarl==1.9.4                                                                                                         
-zipp==1.0.0                                                                                                         
-zope.interface==5.4.0                                                                                               
-zstandard==0.22.0 
+pip install -U "jax[tpu]"==0.4.1 -f https://storage.googleapis.com/jax-relea/ses/libtpu_releases.html
+and run
+https://github.com/infwinston/tpu-example.git