From 09ed0c5d5a32d794ed9f78b70906cbeaff0ef439 Mon Sep 17 00:00:00 2001
From: Ryan Pachauri <rpachauri@google.com>
Date: Thu, 16 Oct 2025 11:39:08 -0700
Subject: [PATCH] Make layer_norm epsilon a constant.

PiperOrigin-RevId: 820327627
Change-Id: Ib7e8ec29db052ac2c4fd79658203bdea2a2845d9
---
 alphafold/model/common_modules.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/alphafold/model/common_modules.py b/alphafold/model/common_modules.py
index 56b0b1f..bac655e 100644
--- a/alphafold/model/common_modules.py
+++ b/alphafold/model/common_modules.py
@@ -25,6 +25,7 @@ import numpy as np
 TRUNCATED_NORMAL_STDDEV_FACTOR = np.asarray(
     0.87962566103423978, dtype=np.float32
 )
+LAYER_NORM_EPSILON = 1e-5
 
 
 def get_initializer_scale(initializer_name, input_shape):
@@ -103,9 +104,6 @@ class Linear(hk.Module):
     Returns:
       output of shape [...] + num_output.
     """
-
-    num_input_dims = self.num_input_dims
-
     if self.num_input_dims > 0:
       in_shape = inputs.shape[-self.num_input_dims :]
     else:
@@ -150,7 +148,7 @@ class LayerNorm(hk.LayerNorm):
       axis,
       create_scale: bool,
       create_offset: bool,
-      eps: float = 1e-5,
+      eps: float = LAYER_NORM_EPSILON,
       scale_init=None,
       offset_init=None,
       use_fast_variance: bool = False,