Merge pull request #67 from AI-Hypercomputer/shangkun-fix-jaxbench

shangkunwang01 · web-flow · commit e73792e5445f · 2026-07-24T11:48:08.000-07:00
fix: update the adapted jaxbench dataset for 19 KernelBench baselines zero-init weights -&gt; random
diff --git a/JAXBench/benchmark/15p_RetNet_Retention/baseline.py b/JAXBench/benchmark/15p_RetNet_Retention/baseline.py
@@ -61,7 +61,7 @@ def workload(query, key, value):
     causal_mask = (distance >= 0).astype(jnp.float32)
     # γ^distance: (H, S, S)
     log_gamma = jnp.log(gammas)  # (H,)
-    decay = jnp.exp(log_gamma[:, None, None] * distance[None, :, :])  # (H, S, S)
+    decay = jnp.exp(log_gamma[:, None, None] * jnp.maximum(distance, 0.0)[None, :, :])  # (H, S, S)
     decay = decay * causal_mask[None, :, :]  # apply causal mask
 
     # Retention: (Q K^T ⊙ D) V
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/18k_Conv2D_ReLU_BiasAdd/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/18k_Conv2D_ReLU_BiasAdd/kernel_task.yaml
@@ -12,11 +12,13 @@ input_gen_code: |-
       height = width = 128
 
       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)
       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
 
       dynamic_args = [x, weight, conv_bias, bias]
       static_args = []
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/18k_Conv2D_ReLU_BiasAdd/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/18k_Conv2D_ReLU_BiasAdd/reference.py
@@ -11,11 +11,13 @@ def get_inputs(dtype=jnp.float32):
     height = width = 128
 
     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)
     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
 
     dynamic_args = [x, weight, conv_bias, bias]
     static_args = []
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/19k_Matmul_Subtract_Multiply_ReLU/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/19k_Matmul_Subtract_Multiply_ReLU/kernel_task.yaml
@@ -12,9 +12,11 @@ input_gen_code: |-
       multiply_value = 1.5
 
       key = jax.random.key(0)
+      rand_key = jax.random.key(0xBADC0DE)
+      ka, kb = jax.random.split(rand_key, 2)
       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
 
       dynamic_args = [x, weight, bias]
       static_args = [subtract_value, multiply_value]
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/19k_Matmul_Subtract_Multiply_ReLU/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/19k_Matmul_Subtract_Multiply_ReLU/reference.py
@@ -11,9 +11,11 @@ def get_inputs(dtype=jnp.float32):
     multiply_value = 1.5
 
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     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
 
     dynamic_args = [x, weight, bias]
     static_args = [subtract_value, multiply_value]
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/20k_Gemm_Multiply_LeakyReLU/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/20k_Gemm_Multiply_LeakyReLU/kernel_task.yaml
@@ -10,9 +10,11 @@ input_gen_code: |-
       out_features = 8192
 
       key = jax.random.key(0)
+      rand_key = jax.random.key(0xBADC0DE)
+      ka, kb = jax.random.split(rand_key, 2)
       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
 
       dynamic_args = [x, weight, bias]
       static_args = []
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/20k_Gemm_Multiply_LeakyReLU/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/20k_Gemm_Multiply_LeakyReLU/reference.py
@@ -9,9 +9,11 @@ def get_inputs(dtype=jnp.float32):
     out_features = 8192
 
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     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
 
     dynamic_args = [x, weight, bias]
     static_args = []
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/22k_Conv2d_InstanceNorm_Divide/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/22k_Conv2d_InstanceNorm_Divide/kernel_task.yaml
@@ -13,10 +13,12 @@ input_gen_code: |-
 
       dtype = jnp.float32
       key = jax.random.key(0)
+      rand_key = jax.random.key(0xBADC0DE)
+      ka, kb = jax.random.split(rand_key, 2)
       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)
 
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/22k_Conv2d_InstanceNorm_Divide/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/22k_Conv2d_InstanceNorm_Divide/reference.py
@@ -12,10 +12,12 @@ def get_inputs():
 
     dtype = jnp.float32
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     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)
 
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/23k_Matmul_Sum_Max_AvgPool_LogSumExp_LogSumExp/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/23k_Matmul_Sum_Max_AvgPool_LogSumExp_LogSumExp/kernel_task.yaml
@@ -9,9 +9,11 @@ input_gen_code: |-
       in_features = 8192
       out_features = 8192
       key = jax.random.key(0)
+      rand_key = jax.random.key(0xBADC0DE)
+      ka, kb = jax.random.split(rand_key, 2)
       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
       dynamic_args = [x, weight, bias]
       static_args = []
       return dynamic_args, static_args
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/23k_Matmul_Sum_Max_AvgPool_LogSumExp_LogSumExp/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/23k_Matmul_Sum_Max_AvgPool_LogSumExp_LogSumExp/reference.py
@@ -8,9 +8,11 @@ def get_inputs(dtype=jnp.float32):
     in_features = 8192
     out_features = 8192
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     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
     dynamic_args = [x, weight, bias]
     static_args = []
     return dynamic_args, static_args
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/27k_Matmul_Mish_Mish/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/27k_Matmul_Mish_Mish/kernel_task.yaml
@@ -9,9 +9,11 @@ input_gen_code: |-
       in_features = 8192
       out_features = 8192
       key = jax.random.key(0)
+      rand_key = jax.random.key(0xBADC0DE)
+      ka, kb = jax.random.split(rand_key, 2)
       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
       dynamic_args = [x, weight, bias]
       static_args = []
       return dynamic_args, static_args
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/27k_Matmul_Mish_Mish/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/27k_Matmul_Mish_Mish/reference.py
@@ -8,9 +8,11 @@ def get_inputs(dtype=jnp.float32):
     in_features = 8192
     out_features = 8192
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     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
     dynamic_args = [x, weight, bias]
     static_args = []
     return dynamic_args, static_args
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/29k_Matmul_Swish_Sum_GroupNorm/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/29k_Matmul_Swish_Sum_GroupNorm/kernel_task.yaml
@@ -6,13 +6,15 @@ input_gen_code: |-
       import jax.numpy as jnp
 
       key = jax.random.key(0)
+      rand_key = jax.random.key(0xBADC0DE)
+      ka, kb = jax.random.split(rand_key, 2)
       batch_size = 8192
       in_features = 4096
       out_features = 4096
       num_groups = 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)
       dynamic_args = [x, weight, bias, gn_weight, gn_bias]
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/29k_Matmul_Swish_Sum_GroupNorm/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/29k_Matmul_Swish_Sum_GroupNorm/reference.py
@@ -5,13 +5,15 @@
 # Initialization
 def get_inputs(dtype=jnp.float32):
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     batch_size = 8192
     in_features = 4096
     out_features = 4096
     num_groups = 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)
     dynamic_args = [x, weight, bias, gn_weight, gn_bias]
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/30k_Matmul_Scaling_ResidualAdd/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/30k_Matmul_Scaling_ResidualAdd/kernel_task.yaml
@@ -9,9 +9,11 @@ input_gen_code: |-
       in_features = 4096
       out_features = 4096
       key = jax.random.key(0)
+      rand_key = jax.random.key(0xBADC0DE)
+      ka, kb = jax.random.split(rand_key, 2)
       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
       dynamic_args = [x, weight, bias]
       static_args = []
       return dynamic_args, static_args
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/30k_Matmul_Scaling_ResidualAdd/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/30k_Matmul_Scaling_ResidualAdd/reference.py
@@ -8,9 +8,11 @@ def get_inputs(dtype=jnp.float32):
     in_features = 4096
     out_features = 4096
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     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
     dynamic_args = [x, weight, bias]
     static_args = []
     return dynamic_args, static_args
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/33k_Conv3d_Mish_Tanh/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/33k_Conv3d_Mish_Tanh/kernel_task.yaml
@@ -12,9 +12,11 @@ input_gen_code: |-
       D, H, W = 32, 64, 64
 
       key = jax.random.key(0)
+      rand_key = jax.random.key(0xBADC0DE)
+      ka, kb = jax.random.split(rand_key, 2)
       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
 
       dynamic_args = [x, weight, bias]
       static_args = []
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/33k_Conv3d_Mish_Tanh/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/33k_Conv3d_Mish_Tanh/reference.py
@@ -11,9 +11,11 @@ def get_inputs(dtype=jnp.float32):
     D, H, W = 32, 64, 64
 
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     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
 
     dynamic_args = [x, weight, bias]
     static_args = []
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/35k_Gemm_Scaling_Hardtanh_GELU/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/35k_Gemm_Scaling_Hardtanh_GELU/kernel_task.yaml
@@ -23,9 +23,11 @@ input_gen_code: |-
       hardtanh_max = CONFIG['hardtanh_max']
 
       key = jax.random.key(0)
+      rand_key = jax.random.key(0xBADC0DE)
+      ka, kb = jax.random.split(rand_key, 2)
       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
 
       dynamic_args = [x, weight, bias]
       static_args = [scaling_factor, hardtanh_min, hardtanh_max]
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/35k_Gemm_Scaling_Hardtanh_GELU/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/35k_Gemm_Scaling_Hardtanh_GELU/reference.py
@@ -22,9 +22,11 @@ def get_inputs(dtype=jnp.float32):
     hardtanh_max = CONFIG['hardtanh_max']
 
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     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
 
     dynamic_args = [x, weight, bias]
     static_args = [scaling_factor, hardtanh_min, hardtanh_max]
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/37k_Matmul_Swish_Scaling/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/37k_Matmul_Swish_Scaling/kernel_task.yaml
@@ -10,9 +10,11 @@ input_gen_code: |-
       out_features = 8192
 
       key = jax.random.key(0)
+      rand_key = jax.random.key(0xBADC0DE)
+      ka, kb = jax.random.split(rand_key, 2)
       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
 
       dynamic_args = [x, weight, bias]
       static_args = []
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/37k_Matmul_Swish_Scaling/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/37k_Matmul_Swish_Scaling/reference.py
@@ -9,9 +9,11 @@ def get_inputs(dtype=jnp.float32):
     out_features = 8192
 
     key = jax.random.key(0)
+    rand_key = jax.random.key(0xBADC0DE)
+    ka, kb = jax.random.split(rand_key, 2)
     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
 
     dynamic_args = [x, weight, bias]
     static_args = []
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/39k_Conv2d_GELU_GlobalAvgPool/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/39k_Conv2d_GELU_GlobalAvgPool/kernel_task.yaml
@@ -6,11 +6,13 @@ input_gen_code: |-
       import jax.numpy as jnp
 
       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=jnp.float32)
-      weight = jnp.zeros((out_channels, in_channels, kernel_size, kernel_size), dtype=jnp.float32)
-      bias = jnp.zeros(out_channels, dtype=jnp.float32)
+      weight = jax.random.normal(ka, (out_channels, in_channels, kernel_size, kernel_size), dtype=jnp.float32) * 0.02
+      bias = jax.random.normal(kb, out_channels, dtype=jnp.float32) * 0.02
       dynamic_args = [x, weight, bias]
       static_args = []
       return dynamic_args, static_args
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/39k_Conv2d_GELU_GlobalAvgPool/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/39k_Conv2d_GELU_GlobalAvgPool/reference.py
@@ -5,11 +5,13 @@
 # Initialization
 def get_inputs():
     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=jnp.float32)
-    weight = jnp.zeros((out_channels, in_channels, kernel_size, kernel_size), dtype=jnp.float32)
-    bias = jnp.zeros(out_channels, dtype=jnp.float32)
+    weight = jax.random.normal(ka, (out_channels, in_channels, kernel_size, kernel_size), dtype=jnp.float32) * 0.02
+    bias = jax.random.normal(kb, out_channels, dtype=jnp.float32) * 0.02
     dynamic_args = [x, weight, bias]
     static_args = []
     return dynamic_args, static_args
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/40k_Gemm_GroupNorm_Min_BiasAdd/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/40k_Gemm_GroupNorm_Min_BiasAdd/kernel_task.yaml
@@ -6,13 +6,15 @@ input_gen_code: |-
       import jax.numpy as jnp
 
       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
       dynamic_args = [x, weight, linear_bias, gn_weight, gn_bias, bias]
       static_args = []
       return dynamic_args, static_args
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/40k_Gemm_GroupNorm_Min_BiasAdd/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/40k_Gemm_GroupNorm_Min_BiasAdd/reference.py
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/41k_Gemm_Add_ReLU/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/41k_Gemm_Add_ReLU/kernel_task.yaml
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/41k_Gemm_Add_ReLU/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/41k_Gemm_Add_ReLU/reference.py
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/44k_Matmul_Divide_GELU/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/44k_Matmul_Divide_GELU/kernel_task.yaml
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/44k_Matmul_Divide_GELU/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/44k_Matmul_Divide_GELU/reference.py
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/45k_Gemm_GroupNorm_Swish_Multiply_Swish/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/45k_Gemm_GroupNorm_Swish_Multiply_Swish/kernel_task.yaml
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/45k_Gemm_GroupNorm_Swish_Multiply_Swish/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/45k_Gemm_GroupNorm_Swish_Multiply_Swish/reference.py
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/47k_Matmul_Add_Swish_Tanh_GELU_Hardtanh/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/47k_Matmul_Add_Swish_Tanh_GELU_Hardtanh/kernel_task.yaml
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/47k_Matmul_Add_Swish_Tanh_GELU_Hardtanh/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/47k_Matmul_Add_Swish_Tanh_GELU_Hardtanh/reference.py
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/48k_Matmul_BatchNorm_BiasAdd_Divide_Swish/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/48k_Matmul_BatchNorm_BiasAdd_Divide_Swish/kernel_task.yaml
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/48k_Matmul_BatchNorm_BiasAdd_Divide_Swish/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/48k_Matmul_BatchNorm_BiasAdd_Divide_Swish/reference.py
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/49k_Matmul_AvgPool_GELU_Scale_Max/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/49k_Matmul_AvgPool_GELU_Scale_Max/kernel_task.yaml
diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/49k_Matmul_AvgPool_GELU_Scale_Max/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/49k_Matmul_AvgPool_GELU_Scale_Max/reference.py
