Ramith Hettiarachchi ramithuh

## jax_default_device.py
import jax
import jax.numpy as jnp

jax.config.update("jax_default_device", jax.devices("cpu")[0])

#runs on cpu
long_vector = jnp.arange(int(1e7))
%timeit jnp.dot(long_vector, long_vector).block_until_ready()

jax.config.update("jax_default_device", jax.devices("gpu")[2])

## log_structure_custom_color_scheme_wandb.ipynb

      
        
          
            
              
              1 file
            
          
          
            
              
              0 forks
            
          
          
            
              
              0 comments
            
          
          
            
              
              0 stars
            
          
        
        
          
              
          
          
            
                ramithuh
                / log_structure_custom_color_scheme_wandb.ipynb
            
            
              Last active
              February 2, 2024 05:16
            
          
        
      
        
  
    
    

          
    
      
      Sorry, something went wrong. Reload?
      Sorry, we cannot display this file.
      Sorry, this file is invalid so it cannot be displayed.
      
          Viewer requires iframe.
	import jax
	import jax.numpy as jnp

	jax.config.update("jax_default_device", jax.devices("cpu")[0])

	#runs on cpu
	long_vector = jnp.arange(int(1e7))
	%timeit jnp.dot(long_vector, long_vector).block_until_ready()

	jax.config.update("jax_default_device", jax.devices("gpu")[2])