Rename use_static_masks -> optimize_for_inference

Author: borzunov

dalle_pytorch/dalle_pytorch.py

@@ -344,7 +344,7 @@ def __init__(
         shared_attn_ids = None,
         shared_ff_ids = None,
         share_input_output_emb = False,
-        use_static_masks = False,
+        optimize_for_inference = False,
     ):
         super().__init__()
         assert isinstance(vae, (DiscreteVAE, OpenAIDiscreteVAE, VQGanVAE)), 'vae must be an instance of DiscreteVAE'
@@ -392,7 +392,7 @@ def __init__(
             rotary_emb = rotary_emb,
             shared_attn_ids = shared_attn_ids,
             shared_ff_ids = shared_ff_ids,
-            use_static_masks = use_static_masks,
+            optimize_for_inference = optimize_for_inference,
         )
 
         self.stable = stable

dalle_pytorch/transformer.py

@@ -220,7 +220,7 @@ def __init__(
         rotary_emb = True,
         shared_attn_ids = None,
         shared_ff_ids = None,
-        use_static_masks = False,
+        optimize_for_inference = False,
     ):
         super().__init__()
         layers = nn.ModuleList([])
@@ -245,12 +245,12 @@ def __init__(
             elif attn_type == 'sparse':
                 attn_class = SparseAttention
             elif attn_type == 'axial_row':
-                if use_static_masks:
+                if optimize_for_inference:
                     attn_class = partial(Attention, stable = stable, static_mask = self._get_static_mask(attn_type))
                 else:
                     attn_class = partial(SparseAxialCausalAttention, seq_len = seq_len, axis = 0, image_size = image_fmap_size, stable = stable)
             elif attn_type == 'axial_col':
-                if use_static_masks:
+                if optimize_for_inference:
                     attn_class = partial(Attention, stable = stable, static_mask = self._get_static_mask(attn_type))
                 else:
                     attn_class = partial(SparseAxialCausalAttention, seq_len = seq_len, axis = 1, image_size = image_fmap_size, stable = stable)