JAXBench: Fix no-op exploit in 19 KernelBench baselines: zero-init weights -> random

JAXBench/benchmark/18k_Conv2D_ReLU_BiasAdd/baseline.py

@@ -14,13 +14,15 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb, kc = jax.random.split(rand_key, 3)
     k1, k2 = jax.random.split(key)
     batch_size, in_channels, out_channels, kernel_size = 128, 64, 128, 3
     height = width = 128
     x = jax.random.uniform(k1, (batch_size, in_channels, height, width), dtype=dtype)
-    weight = jnp.zeros((out_channels, in_channels, kernel_size, kernel_size), dtype=dtype)
-    conv_bias = jnp.zeros(out_channels, dtype=dtype)
-    bias = jnp.zeros((out_channels, 1, 1), dtype=dtype)
+    weight = jax.random.normal(ka, (out_channels, in_channels, kernel_size, kernel_size), dtype=dtype) * 0.02
+    conv_bias = jax.random.normal(kb, out_channels, dtype=dtype) * 0.02
+    bias = jax.random.normal(kc, (out_channels, 1, 1), dtype=dtype) * 0.02
     return x, weight, conv_bias, bias
 
 

JAXBench/benchmark/19k_Matmul_Subtract_Multiply_ReLU/baseline.py

@@ -15,9 +15,11 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     x = jax.random.uniform(key, (4096, 8192), dtype=dtype)
-    weight = jnp.zeros((8192, 8192), dtype=dtype)
-    bias = jnp.zeros(8192, dtype=dtype)
+    weight = jax.random.normal(ka, (8192, 8192), dtype=dtype) * 0.02
+    bias = jax.random.normal(kb, 8192, dtype=dtype) * 0.02
     return x, weight, bias
 
 

JAXBench/benchmark/20k_Gemm_Multiply_LeakyReLU/baseline.py

@@ -15,9 +15,11 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     x = jax.random.uniform(key, (4096, 8192), dtype=dtype)
-    weight = jnp.zeros((8192, 8192), dtype=dtype)
-    bias = jnp.zeros(8192, dtype=dtype)
+    weight = jax.random.normal(ka, (8192, 8192), dtype=dtype) * 0.02
+    bias = jax.random.normal(kb, 8192, dtype=dtype) * 0.02
     return x, weight, bias
 
 

JAXBench/benchmark/22k_Conv2d_InstanceNorm_Divide/baseline.py

@@ -15,11 +15,13 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     batch_size, in_channels, out_channels, kernel_size = 128, 64, 128, 3
     height = width = 128
     x = jax.random.uniform(key, (batch_size, in_channels, height, width), dtype=dtype)
-    weight = jnp.zeros((out_channels, in_channels, kernel_size, kernel_size), dtype=dtype)
-    conv_bias = jnp.zeros(out_channels, dtype=dtype)
+    weight = jax.random.normal(ka, (out_channels, in_channels, kernel_size, kernel_size), dtype=dtype) * 0.02
+    conv_bias = jax.random.normal(kb, out_channels, dtype=dtype) * 0.02
     in_weight = jnp.ones(out_channels, dtype=dtype)
     in_bias = jnp.zeros(out_channels, dtype=dtype)
     return x, weight, conv_bias, in_weight, in_bias

JAXBench/benchmark/23k_Matmul_Sum_Max_AvgPool_LogSumExp_LogSumExp/baseline.py

@@ -13,9 +13,11 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     x = jax.random.uniform(key, (4096, 8192), dtype=dtype)
-    weight = jnp.zeros((8192, 8192), dtype=dtype)
-    bias = jnp.zeros(8192, dtype=dtype)
+    weight = jax.random.normal(ka, (8192, 8192), dtype=dtype) * 0.02
+    bias = jax.random.normal(kb, 8192, dtype=dtype) * 0.02
     return x, weight, bias
 
 

JAXBench/benchmark/27k_Matmul_Mish_Mish/baseline.py

@@ -13,9 +13,11 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     x = jax.random.uniform(key, (4096, 8192), dtype=dtype)
-    weight = jnp.zeros((8192, 8192), dtype=dtype)
-    bias = jnp.zeros(8192, dtype=dtype)
+    weight = jax.random.normal(ka, (8192, 8192), dtype=dtype) * 0.02
+    bias = jax.random.normal(kb, 8192, dtype=dtype) * 0.02
     return x, weight, bias
 
 

JAXBench/benchmark/29k_Matmul_Swish_Sum_GroupNorm/baseline.py

@@ -14,10 +14,12 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     batch_size, in_features, out_features, num_groups = 8192, 4096, 4096, 64
     x = jax.random.uniform(key, (batch_size, in_features), dtype=dtype)
-    weight = jnp.zeros((in_features, out_features), dtype=dtype)
-    bias = jnp.zeros(out_features, dtype=dtype)
+    weight = jax.random.normal(ka, (in_features, out_features), dtype=dtype) * 0.02
+    bias = jax.random.normal(kb, out_features, dtype=dtype) * 0.02
     gn_weight = jnp.ones(out_features, dtype=dtype)
     gn_bias = jnp.zeros(out_features, dtype=dtype)
     return x, weight, bias, gn_weight, gn_bias

JAXBench/benchmark/30k_Matmul_Scaling_ResidualAdd/baseline.py

@@ -14,9 +14,11 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     x = jax.random.uniform(key, (16384, 4096), dtype=dtype)
-    weight = jnp.zeros((4096, 4096), dtype=dtype)
-    bias = jnp.zeros(4096, dtype=dtype)
+    weight = jax.random.normal(ka, (4096, 4096), dtype=dtype) * 0.02
+    bias = jax.random.normal(kb, 4096, dtype=dtype) * 0.02
     return x, weight, bias
 
 

JAXBench/benchmark/33k_Conv3d_Mish_Tanh/baseline.py

@@ -14,11 +14,13 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     batch_size, in_channels, out_channels, kernel_size = 16, 32, 64, 3
     D, H, W = 32, 64, 64
     x = jax.random.uniform(key, (batch_size, in_channels, D, H, W), dtype=dtype)
-    weight = jnp.zeros((out_channels, in_channels, kernel_size, kernel_size, kernel_size), dtype=dtype)
-    bias = jnp.zeros(out_channels, dtype=dtype)
+    weight = jax.random.normal(ka, (out_channels, in_channels, kernel_size, kernel_size, kernel_size), dtype=dtype) * 0.02
+    bias = jax.random.normal(kb, out_channels, dtype=dtype) * 0.02
     return x, weight, bias
 
 

JAXBench/benchmark/35k_Gemm_Scaling_Hardtanh_GELU/baseline.py

@@ -16,9 +16,11 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     x = jax.random.uniform(key, (4096, 8192), dtype=dtype)
-    weight = jnp.zeros((8192, 8192), dtype=dtype)
-    bias = jnp.zeros(8192, dtype=dtype)
+    weight = jax.random.normal(ka, (8192, 8192), dtype=dtype) * 0.02
+    bias = jax.random.normal(kb, 8192, dtype=dtype) * 0.02
     return x, weight, bias
 
 

JAXBench/benchmark/37k_Matmul_Swish_Scaling/baseline.py

@@ -14,9 +14,11 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     x = jax.random.uniform(key, (4096, 8192), dtype=dtype)
-    weight = jnp.zeros((8192, 8192), dtype=dtype)
-    bias = jnp.zeros(8192, dtype=dtype)
+    weight = jax.random.normal(ka, (8192, 8192), dtype=dtype) * 0.02
+    bias = jax.random.normal(kb, 8192, dtype=dtype) * 0.02
     return x, weight, bias
 
 

JAXBench/benchmark/39k_Conv2d_GELU_GlobalAvgPool/baseline.py

@@ -14,11 +14,13 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     batch_size, in_channels, out_channels, kernel_size = 128, 8, 64, 3
     height, width = 256, 256
     x = jax.random.uniform(key, (batch_size, in_channels, height, width), dtype=dtype)
-    weight = jnp.zeros((out_channels, in_channels, kernel_size, kernel_size), dtype=dtype)
-    bias = jnp.zeros(out_channels, dtype=dtype)
+    weight = jax.random.normal(ka, (out_channels, in_channels, kernel_size, kernel_size), dtype=dtype) * 0.02
+    bias = jax.random.normal(kb, out_channels, dtype=dtype) * 0.02
     return x, weight, bias
 
 

JAXBench/benchmark/40k_Gemm_GroupNorm_Min_BiasAdd/baseline.py

@@ -14,13 +14,15 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb, kc = jax.random.split(rand_key, 3)
     batch_size, in_features, out_features, num_groups = 4096, 8192, 8192, 512
     x = jax.random.uniform(key, (batch_size, in_features), dtype=dtype)
-    weight = jnp.zeros((out_features, in_features), dtype=dtype)
-    linear_bias = jnp.zeros(out_features, dtype=dtype)
+    weight = jax.random.normal(ka, (out_features, in_features), dtype=dtype) * 0.02
+    linear_bias = jax.random.normal(kb, out_features, dtype=dtype) * 0.02
     gn_weight = jnp.ones(out_features, dtype=dtype)
     gn_bias = jnp.zeros(out_features, dtype=dtype)
-    bias = jnp.zeros((1, out_features, 1, 1), dtype=dtype)
+    bias = jax.random.normal(kc, (1, out_features, 1, 1), dtype=dtype) * 0.02
     return x, weight, linear_bias, gn_weight, gn_bias, bias
 
 

JAXBench/benchmark/41k_Gemm_Add_ReLU/baseline.py

@@ -13,9 +13,11 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     x = jax.random.uniform(key, (4096, 8192), dtype=dtype)
-    weight = jnp.zeros((8192, 8192), dtype=dtype)
-    bias = jnp.zeros(8192, dtype=dtype)
+    weight = jax.random.normal(ka, (8192, 8192), dtype=dtype) * 0.02
+    bias = jax.random.normal(kb, 8192, dtype=dtype) * 0.02
     return x, weight, bias
 
 

JAXBench/benchmark/44k_Matmul_Divide_GELU/baseline.py

@@ -14,9 +14,11 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     x = jax.random.uniform(key, (4096, 8192), dtype=dtype)
-    weight = jnp.zeros((8192, 8192), dtype=dtype)
-    bias = jnp.zeros(8192, dtype=dtype)
+    weight = jax.random.normal(ka, (8192, 8192), dtype=dtype) * 0.02
+    bias = jax.random.normal(kb, 8192, dtype=dtype) * 0.02
     return x, weight, bias
 
 

JAXBench/benchmark/45k_Gemm_GroupNorm_Swish_Multiply_Swish/baseline.py

@@ -14,13 +14,15 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb, kc = jax.random.split(rand_key, 3)
     batch_size, in_features, out_features, num_groups = 4096, 8192, 8192, 256
     x = jax.random.uniform(key, (batch_size, in_features), dtype=dtype)
-    gemm_weight = jnp.zeros((out_features, in_features), dtype=dtype)
-    gemm_bias = jnp.zeros(out_features, dtype=dtype)
+    gemm_weight = jax.random.normal(ka, (out_features, in_features), dtype=dtype) * 0.02
+    gemm_bias = jax.random.normal(kb, out_features, dtype=dtype) * 0.02
     gn_weight = jnp.ones(out_features, dtype=dtype)
     gn_bias = jnp.zeros(out_features, dtype=dtype)
-    multiply_weight = jnp.zeros(out_features, dtype=dtype)
+    multiply_weight = jax.random.normal(kc, out_features, dtype=dtype) * 0.02
     return x, gemm_weight, gemm_bias, gn_weight, gn_bias, multiply_weight
 
 

JAXBench/benchmark/47k_Matmul_Add_Swish_Tanh_GELU_Hardtanh/baseline.py

@@ -13,10 +13,12 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb, kc = jax.random.split(rand_key, 3)
     x = jax.random.uniform(key, (4096, 8192), dtype=dtype)
-    weight = jnp.zeros((8192, 8192), dtype=dtype)
-    bias = jnp.zeros(8192, dtype=dtype)
-    add_value = jnp.zeros(8192, dtype=dtype)
+    weight = jax.random.normal(ka, (8192, 8192), dtype=dtype) * 0.02
+    bias = jax.random.normal(kb, 8192, dtype=dtype) * 0.02
+    add_value = jax.random.normal(kc, 8192, dtype=dtype) * 0.02
     return x, weight, bias, add_value
 
 

JAXBench/benchmark/48k_Matmul_BatchNorm_BiasAdd_Divide_Swish/baseline.py

@@ -16,15 +16,17 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb, kc = jax.random.split(rand_key, 3)
     batch_size, in_features, out_features = 4096, 8192, 8192
     x = jax.random.uniform(key, (batch_size, in_features), dtype=dtype)
-    weight = jnp.zeros((in_features, out_features), dtype=dtype)
-    linear_bias = jnp.zeros(out_features, dtype=dtype)
+    weight = jax.random.normal(ka, (in_features, out_features), dtype=dtype) * 0.02
+    linear_bias = jax.random.normal(kb, out_features, dtype=dtype) * 0.02
     bn_scale = jnp.ones(out_features, dtype=dtype)
     bn_bias = jnp.zeros(out_features, dtype=dtype)
     bn_mean = jnp.zeros(out_features, dtype=dtype)
     bn_var = jnp.ones(out_features, dtype=dtype)
-    bias = jnp.zeros((1,), dtype=dtype)
+    bias = jax.random.normal(kc, (1,), dtype=dtype) * 0.02
     return x, weight, linear_bias, bn_scale, bn_bias, bn_mean, bn_var, bias
 
 

JAXBench/benchmark/49k_Matmul_AvgPool_GELU_Scale_Max/baseline.py

@@ -15,9 +15,11 @@
 def create_inputs(dtype=jnp.float32):
     """Create all inputs including weights."""
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     x = jax.random.uniform(key, (4096, 8192), dtype=dtype)
-    weight = jnp.zeros((8192, 8192), dtype=dtype)
-    bias = jnp.zeros(8192, dtype=dtype)
+    weight = jax.random.normal(ka, (8192, 8192), dtype=dtype) * 0.02
+    bias = jax.random.normal(kb, 8192, dtype=dtype) * 0.02
     return x, weight, bias