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

18k-50k KernelBench-derived baselines initialized their weight/bias tensors to
jnp.zeros in create_inputs. With zero weights, `x @ W (+ b)` is identically
zero and independent of the input, so the reference output is a trivial
constant (all-zero, or a fixed activation thereof). Any kernel returning that
constant -- including a no-op that skips the matmul/conv entirely -- passes
np.allclose, so these benchmarks could report large meaningless speedups
without computing the operator.

This replaces the zero-init weights/biases with small-normal random values
(~0.02 scale: input-dependent, bf16-representable, no overflow). Only
create_inputs is changed; the workload/op is untouched. After the fix a no-op
(all-zero output) fails correctness on all 19.

Scope: 19 of the affected baselines are fully fixed by non-zero weights. Five
others whose *output* is intrinsically small regardless of weights -- the
softmax-terminated 38k/43k/50k (row outputs ~1/N) and the structurally
degenerate 25k (GroupNorm->Mean) and 42k (Max-Subtract-GELU) -- are NOT
addressed here; they need a tolerance or operator change and are left to a
follow-up. Megablox (11p) has a distinct input-underflow variant fixed
separately.

Author: charleshong3

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