From c7991dee856ea2c7df744b1403729f866cd3eb7f Mon Sep 17 00:00:00 2001
From: Jean Charle Yaacoub <50300488+jyaacoub@users.noreply.github.com>
Date: Mon, 29 Apr 2024 15:29:50 -0400
Subject: [PATCH] Environment vars for LMA default chunk_sizes

This should be a config arg but I am not sure how to go about that without messing things up.
---
 openfold/model/primitives.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/openfold/model/primitives.py b/openfold/model/primitives.py
index ea38cb34a..31a36df9f 100644
--- a/openfold/model/primitives.py
+++ b/openfold/model/primitives.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import importlib
-import math
+import math, os
 from typing import Optional, Callable, List, Tuple
 import numpy as np
 
@@ -42,9 +42,8 @@
     flatten_final_dims,
 )
 
-
-DEFAULT_LMA_Q_CHUNK_SIZE = 1024
-DEFAULT_LMA_KV_CHUNK_SIZE = 4096
+DEFAULT_LMA_Q_CHUNK_SIZE = int(os.environ.get('DEFAULT_LMA_Q_CHUNK_SIZE', 1024))
+DEFAULT_LMA_KV_CHUNK_SIZE = int(os.environ.get('DEFAULT_LMA_KV_CHUNK_SIZE', 4096))
 
 
 def _prod(nums):