p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] test/test_jax.py (python)
# import jax
# with jax.profiler.trace("tensorboard"):
  # key = jax.random.key(0)
  # x = jax.random.normal(key, (1024, 1024))
  # y = x @ x
  # y.block_until_ready()

import datetime
import jax.numpy as jnp
import jax

MATRIX_DIM = 32768
STEPS = 10

A = jnp.ones((MATRIX_DIM, MATRIX_DIM), dtype=jnp.float32)
B = jnp.ones((MATRIX_DIM, MATRIX_DIM), dtype=jnp.float32)

num_bytes = A.size * 4
total_num_bytes_crossing_to_hbm = num_bytes * 3

total_num_flops = 2 * MATRIX_DIM * MATRIX_DIM**2

def matmul(A, B):
  return A @ B

matmul(A, B) # warmup

with jax.profiler.trace("tensorboard"):
  start_time = datetime.datetime.now()
  for i in range(STEPS):
    C = A @ B
    C.block_until_ready()
  end_time = datetime.datetime.now()

average_time_per_step = (end_time - start_time).total_seconds() / STEPS

print(f"{average_time_per_step}, teraflops per second: {total_num_flops / average_time_per_step / 1e12}, gigabytes per second: {total_num_bytes_crossing_to_hbm / average_time_per_step / 1e9}")