Juan Acevedo entrpn
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| """ | |
| A bare bones examples of optimizing a black-box function (f) using | |
| Natural Evolution Strategies (NES), where the parameter distribution is a | |
| gaussian of fixed standard deviation. | |
| """ | |
| import numpy as np | |
| np.random.seed(0) | |
| # the function we want to optimize |
jbochi
/ evaluation.ipynb
Last active
May 3, 2025 05:43
Recommending GitHub repositories with Google Big Query and implicit library: https://medium.com/@jbochi/recommending-github-repositories-with-google-bigquery-and-the-implicit-library-e6cce666c77
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
skye
/ tpu_topology_env_vars.py
Last active
June 17, 2025 12:51
You can use these environment variables to run a Python process on a subset of the TPU cores on a Cloud TPU VM. This allows running multiple TPU processes at the same time, since only one process can access a given TPU chip at a time. Note that on TPU v2 and v3, 1 TPU chip = 2 TpuDevice as reported by `jax.devices()` (8 devices total). On v4, 1 …
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| # ==== Non-communicating processes | |
| # 4x 1 chip per process: | |
| os.environ["TPU_CHIPS_PER_PROCESS_BOUNDS"] = "1,1,1" | |
| os.environ["TPU_PROCESS_BOUNDS"] = "1,1,1" | |
| # Different per process: | |
| os.environ["TPU_VISIBLE_DEVICES"] = "0" # "1", "2", "3" | |
| # 1-liner for bash: TPU_CHIPS_PER_PROCESS_BOUNDS=1,1,1 TPU_PROCESS_BOUNDS=1,1,1 TPU_VISIBLE_DEVICES=0 TPU_MESH_CONTROLLER_ADDRESS=localhost:8476 TPU_MESH_CONTROLLER_PORT=8476 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| def forward(self, x, timesteps=None, context=None, y=None, **kwargs): | |
| # broadcast timesteps to batch dimension | |
| timesteps = timesteps.expand(x.shape[0]) | |
| hs = [] | |
| t_emb = get_timestep_embedding(timesteps, self.model_channels) # , repeat_only=False) | |
| t_emb = t_emb.to(x.dtype) | |
| emb = self.time_embed(t_emb) | |
| assert x.shape[0] == y.shape[0], f"batch size mismatch: {x.shape[0]} != {y.shape[0]}" |