⚡️ Speed up method FluxControlLoRALayer.get_parameters by 7%

[codeflash-ai]

📄 7% (0.07x) speedup for FluxControlLoRALayer.get_parameters in invokeai/backend/patches/layers/flux_control_lora_layer.py

⏱️ Runtime : 1.55 milliseconds → 1.45 milliseconds (best of 130 runs)

📝 Explanation and details

The optimized code achieves a 7% speedup by eliminating unnecessary tensor reshaping operations in the get_weight method of LoRALayer.

Key optimizations:

1. Conditional reshaping: Instead of always calling reshape() on tensors, the code first checks if tensors are already 2D using tensor.dim() == 2. This avoids redundant reshape operations when tensors are already in the correct format.

2. Separate variable assignment: The optimized version assigns reshaped tensors to separate variables (up_reshaped, down_reshaped) rather than inlining the reshape operations in the matrix multiplication. This reduces the computational overhead of the @ operator by ensuring it operates on pre-processed tensors.

Performance impact:

• The line profiler shows the most significant improvement in the matrix multiplication line (from 401,940ns to 235,222ns per hit - a 41% reduction)
• Reshape operations are reduced from ~48,000ns total to ~16,500ns total across both up and down tensors
• The optimization is particularly effective for scenarios where tensors are already 2D, as evidenced by test cases showing 10-36% improvements

Real-world benefits:
This optimization is valuable for LoRA (Low-Rank Adaptation) layers commonly used in AI model fine-tuning, where get_weight is called frequently during forward passes. The conditional reshaping reduces computational overhead without changing the mathematical correctness, making it especially beneficial for models with many LoRA layers or during batch processing scenarios.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 152 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import pytest  # used for our unit tests
import torch
from invokeai.backend.patches.layers.flux_control_lora_layer import \
    FluxControlLoRALayer


# Helper to create a dummy LoRA layer with controlled behavior
class DummyFluxControlLoRALayer(FluxControlLoRALayer):
    def __init__(self, up, down, alpha, bias, rank):
        super().__init__(up, None, down, alpha, bias)
        self._test_rank = rank

    def _rank(self):
        return self._test_rank

    def get_weight(self, orig_weight):
        # For testing, just return orig_weight + 1
        return orig_weight + 1

    def get_bias(self, orig_bias):
        # For testing, just return self.bias
        return self.bias

# ----------- UNIT TESTS ------------

# ----------- BASIC TEST CASES ------------

def test_basic_weight_and_bias_present():
    """Basic: Both weight and bias present, alpha and rank set, scale applies."""
    layer = DummyFluxControlLoRALayer(
        up=torch.tensor([[1.0]]), down=torch.tensor([[1.0]]), alpha=2.0, bias=torch.tensor([5.0]), rank=2
    )
    orig_params = {"weight": torch.tensor([10.0]), "bias": torch.tensor([20.0])}
    weight = 3.0
    codeflash_output = layer.get_parameters(orig_params, weight); out = codeflash_output # 39.4μs -> 38.5μs (2.34% faster)

def test_basic_weight_only():
    """Basic: Only weight present, no bias in orig_parameters nor in layer."""
    layer = DummyFluxControlLoRALayer(
        up=torch.tensor([[1.0]]), down=torch.tensor([[1.0]]), alpha=1.0, bias=None, rank=1
    )
    orig_params = {"weight": torch.tensor([2.0])}
    weight = 4.0
    codeflash_output = layer.get_parameters(orig_params, weight); out = codeflash_output # 25.4μs -> 24.5μs (3.66% faster)

def test_basic_weight_and_bias_none_in_orig():
    """Basic: Bias present in layer, but not in orig_parameters dict."""
    layer = DummyFluxControlLoRALayer(
        up=torch.tensor([[1.0]]), down=torch.tensor([[1.0]]), alpha=2.0, bias=torch.tensor([7.0]), rank=2
    )
    orig_params = {"weight": torch.tensor([3.0])}  # no bias key
    weight = 2.0
    codeflash_output = layer.get_parameters(orig_params, weight); out = codeflash_output # 25.6μs -> 24.7μs (3.62% faster)

def test_basic_scale_default_if_alpha_none():
    """Basic: If alpha is None, scale is 1.0."""
    layer = DummyFluxControlLoRALayer(
        up=torch.tensor([[1.0]]), down=torch.tensor([[1.0]]), alpha=None, bias=torch.tensor([2.0]), rank=5
    )
    orig_params = {"weight": torch.tensor([4.0]), "bias": torch.tensor([8.0])}
    weight = 2.0
    codeflash_output = layer.get_parameters(orig_params, weight); out = codeflash_output # 24.7μs -> 24.5μs (0.587% faster)

def test_basic_scale_default_if_rank_none():
    """Basic: If rank is None, scale is 1.0."""
    layer = DummyFluxControlLoRALayer(
        up=torch.tensor([[1.0]]), down=torch.tensor([[1.0]]), alpha=3.0, bias=torch.tensor([6.0]), rank=None
    )
    orig_params = {"weight": torch.tensor([2.0]), "bias": torch.tensor([5.0])}
    weight = 3.0
    codeflash_output = layer.get_parameters(orig_params, weight); out = codeflash_output # 25.0μs -> 24.1μs (3.59% faster)

# ----------- EDGE TEST CASES ------------

def test_edge_weight_zero():
    """Edge: weight argument is zero, output should be zero tensors."""
    layer = DummyFluxControlLoRALayer(
        up=torch.tensor([[1.0]]), down=torch.tensor([[1.0]]), alpha=1.0, bias=torch.tensor([3.0]), rank=1
    )
    orig_params = {"weight": torch.tensor([5.0]), "bias": torch.tensor([7.0])}
    weight = 0.0
    codeflash_output = layer.get_parameters(orig_params, weight); out = codeflash_output # 25.1μs -> 24.2μs (3.81% faster)

def test_edge_negative_weight():
    """Edge: weight argument is negative, output should be negative tensors."""
    layer = DummyFluxControlLoRALayer(
        up=torch.tensor([[1.0]]), down=torch.tensor([[1.0]]), alpha=2.0, bias=torch.tensor([-1.0]), rank=2
    )
    orig_params = {"weight": torch.tensor([-3.0]), "bias": torch.tensor([-2.0])}
    weight = -2.0
    codeflash_output = layer.get_parameters(orig_params, weight); out = codeflash_output # 25.3μs -> 24.3μs (3.89% faster)

def test_edge_bias_is_none_in_layer_and_orig():
    """Edge: No bias in layer nor in orig_parameters."""
    layer = DummyFluxControlLoRALayer(
        up=torch.tensor([[1.0]]), down=torch.tensor([[1.0]]), alpha=1.0, bias=None, rank=1
    )
    orig_params = {"weight": torch.tensor([1.0])}
    weight = 1.0
    codeflash_output = layer.get_parameters(orig_params, weight); out = codeflash_output # 21.5μs -> 21.1μs (2.00% faster)

def test_edge_weight_tensor_shape():
    """Edge: weight tensor is multi-dimensional."""
    layer = DummyFluxControlLoRALayer(
        up=torch.tensor([[1.0, 2.0], [3.0, 4.0]]), down=torch.tensor([[1.0, 2.0], [3.0, 4.0]]), alpha=2.0, bias=torch.tensor([[5.0, 6.0], [7.0, 8.0]]), rank=2
    )
    orig_params = {"weight": torch.tensor([[1.0, 2.0], [3.0, 4.0]])}
    weight = 2.0
    codeflash_output = layer.get_parameters(orig_params, weight); out = codeflash_output # 26.0μs -> 26.1μs (0.314% slower)
    # get_weight = orig + 1, times 2
    expected = (orig_params["weight"] + 1) * 2

def test_edge_weight_tensor_empty():
    """Edge: weight tensor is empty."""
    layer = DummyFluxControlLoRALayer(
        up=torch.tensor([[1.0]]), down=torch.tensor([[1.0]]), alpha=1.0, bias=None, rank=1
    )
    orig_params = {"weight": torch.tensor([])}
    weight = 1.0
    codeflash_output = layer.get_parameters(orig_params, weight); out = codeflash_output # 20.4μs -> 19.2μs (6.01% faster)

def test_edge_bias_tensor_empty():
    """Edge: bias tensor is empty."""
    layer = DummyFluxControlLoRALayer(
        up=torch.tensor([[1.0]]), down=torch.tensor([[1.0]]), alpha=1.0, bias=torch.tensor([]), rank=1
    )
    orig_params = {"weight": torch.tensor([1.0]), "bias": torch.tensor([])}
    weight = 1.0
    codeflash_output = layer.get_parameters(orig_params, weight); out = codeflash_output # 24.5μs -> 25.1μs (2.55% slower)

def test_edge_weight_tensor_dtype_float32():
    """Edge: weight tensor is float32, output should preserve dtype."""
    layer = DummyFluxControlLoRALayer(
        up=torch.tensor([[1.0]], dtype=torch.float32), down=torch.tensor([[1.0]], dtype=torch.float32), alpha=1.0, bias=torch.tensor([2.0], dtype=torch.float32), rank=1
    )
    orig_params = {"weight": torch.tensor([2.0], dtype=torch.float32)}
    weight = 2.0
    codeflash_output = layer.get_parameters(orig_params, weight); out = codeflash_output # 25.1μs -> 24.6μs (1.93% faster)

def test_edge_weight_tensor_dtype_int():
    """Edge: weight tensor is int, output should preserve dtype."""
    layer = DummyFluxControlLoRALayer(
        up=torch.tensor([[1]], dtype=torch.int64), down=torch.tensor([[1]], dtype=torch.int64), alpha=1, bias=torch.tensor([2], dtype=torch.int64), rank=1
    )
    orig_params = {"weight": torch.tensor([2], dtype=torch.int64)}
    weight = 2
    codeflash_output = layer.get_parameters(orig_params, weight); out = codeflash_output # 27.1μs -> 26.6μs (1.55% faster)

def test_edge_weight_tensor_device_cpu():
    """Edge: weight tensor is on CPU."""
    layer = DummyFluxControlLoRALayer(
        up=torch.tensor([[1.0]]), down=torch.tensor([[1.0]]), alpha=1.0, bias=torch.tensor([2.0]), rank=1
    )
    orig_params = {"weight": torch.tensor([2.0])}
    weight = 2.0
    codeflash_output = layer.get_parameters(orig_params, weight); out = codeflash_output # 24.9μs -> 24.1μs (3.49% faster)

@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")

def test_edge_missing_weight_key_raises():
    """Edge: orig_parameters missing 'weight' key should raise KeyError."""
    layer = DummyFluxControlLoRALayer(
        up=torch.tensor([[1.0]]), down=torch.tensor([[1.0]]), alpha=1.0, bias=None, rank=1
    )
    orig_params = {"bias": torch.tensor([1.0])}
    weight = 1.0
    with pytest.raises(KeyError):
        layer.get_parameters(orig_params, weight) # 2.37μs -> 2.35μs (0.853% faster)

# ----------- LARGE SCALE TEST CASES ------------

def test_large_scale_weight_tensor_1000_elements():
    """Large scale: weight tensor with 1000 elements."""
    layer = DummyFluxControlLoRALayer(
        up=torch.ones((10, 10)), down=torch.ones((10, 10)), alpha=10.0, bias=torch.ones(1000), rank=10
    )
    orig_params = {"weight": torch.ones(1000), "bias": torch.ones(1000)}
    weight = 2.0
    codeflash_output = layer.get_parameters(orig_params, weight); out = codeflash_output # 36.3μs -> 35.3μs (2.98% faster)

def test_large_scale_weight_tensor_2d():
    """Large scale: weight tensor is 100x10 matrix."""
    layer = DummyFluxControlLoRALayer(
        up=torch.ones((10, 10)), down=torch.ones((10, 10)), alpha=10.0, bias=torch.ones((100, 10)), rank=10
    )
    orig_params = {"weight": torch.ones((100, 10)), "bias": torch.ones((100, 10))}
    weight = 3.0
    codeflash_output = layer.get_parameters(orig_params, weight); out = codeflash_output # 26.7μs -> 25.9μs (3.11% faster)

def test_large_scale_weight_tensor_3d():
    """Large scale: weight tensor is 10x10x10 tensor."""
    layer = DummyFluxControlLoRALayer(
        up=torch.ones((10, 10)), down=torch.ones((10, 10)), alpha=10.0, bias=torch.ones((10, 10, 10)), rank=10
    )
    orig_params = {"weight": torch.ones((10, 10, 10)), "bias": torch.ones((10, 10, 10))}
    weight = 1.5
    codeflash_output = layer.get_parameters(orig_params, weight); out = codeflash_output # 26.1μs -> 24.8μs (5.05% faster)

def test_large_scale_weight_tensor_max_size():
    """Large scale: weight tensor with 1000 elements, float32, total size < 100MB."""
    layer = DummyFluxControlLoRALayer(
        up=torch.ones((10, 10)), down=torch.ones((10, 10)), alpha=10.0, bias=torch.ones(1000, dtype=torch.float32), rank=10
    )
    orig_params = {"weight": torch.ones(1000, dtype=torch.float32), "bias": torch.ones(1000, dtype=torch.float32)}
    weight = 1.0
    codeflash_output = layer.get_parameters(orig_params, weight); out = codeflash_output # 24.4μs -> 23.8μs (2.56% faster)

def test_large_scale_weight_tensor_no_bias():
    """Large scale: weight tensor with 1000 elements, no bias."""
    layer = DummyFluxControlLoRALayer(
        up=torch.ones((10, 10)), down=torch.ones((10, 10)), alpha=10.0, bias=None, rank=10
    )
    orig_params = {"weight": torch.ones(1000)}
    weight = 2.0
    codeflash_output = layer.get_parameters(orig_params, weight); out = codeflash_output # 22.1μs -> 20.0μs (10.1% faster)

def test_large_scale_weight_tensor_random_values():
    """Large scale: weight tensor with random values, check computation."""
    orig_weight = torch.arange(1000, dtype=torch.float32)
    orig_bias = torch.arange(1000, dtype=torch.float32)
    layer = DummyFluxControlLoRALayer(
        up=torch.ones((10, 10)), down=torch.ones((10, 10)), alpha=10.0, bias=orig_bias, rank=10
    )
    orig_params = {"weight": orig_weight, "bias": orig_bias}
    weight = 2.0
    codeflash_output = layer.get_parameters(orig_params, weight); out = codeflash_output # 24.7μs -> 24.1μs (2.60% faster)
    # get_weight = orig_weight+1, times 2
    expected_weight = (orig_weight + 1) * 2
    expected_bias = orig_bias * 2
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import pytest  # used for our unit tests
import torch
from invokeai.backend.patches.layers.flux_control_lora_layer import \
    FluxControlLoRALayer

# ---- UNIT TESTS ----

# Basic Test Cases

def test_basic_weight_and_bias():
    # Test with simple up/down matrices, scalar alpha, and bias
    up = torch.ones(2, 3)
    down = torch.ones(3, 4)
    bias = torch.ones(4)
    orig_weight = torch.ones(2, 4)
    orig_bias = torch.ones(4)
    alpha = 6.0
    # rank = 3 (up.shape[1])
    lora = FluxControlLoRALayer(up, None, down, alpha, bias)
    codeflash_output = lora.get_parameters({"weight": orig_weight, "bias": orig_bias}, weight=2.0); params = codeflash_output # 43.9μs -> 32.9μs (33.2% faster)
    # scale = alpha / rank = 2.0
    # weight = up @ down * (weight * scale) = (2x3 @ 3x4) = 2x4, all values = 3 (since all ones)
    expected_weight = torch.full((2, 4), 3.0) * (2.0 * 2.0)
    expected_bias = torch.ones(4) * (2.0 * 2.0)

def test_basic_weight_no_bias():
    # Test with bias=None, should not return bias in params
    up = torch.ones(2, 2)
    down = torch.ones(2, 2)
    orig_weight = torch.ones(2, 2)
    alpha = 4.0
    lora = FluxControlLoRALayer(up, None, down, alpha, None)
    codeflash_output = lora.get_parameters({"weight": orig_weight}, weight=1.5); params = codeflash_output # 31.5μs -> 23.5μs (34.0% faster)
    # scale = alpha / rank = 2.0
    expected_weight = torch.full((2, 2), 2.0) * (1.5 * 2.0)

def test_basic_mid_tensor():
    # Test with mid tensor present
    up = torch.ones(2, 2)
    down = torch.ones(2, 2)
    mid = torch.ones(2, 2, 1, 1)
    orig_weight = torch.ones(2, 2, 1, 1)
    alpha = 2.0
    lora = FluxControlLoRALayer(up, mid, down, alpha, None)
    codeflash_output = lora.get_parameters({"weight": orig_weight}, weight=1.0); params = codeflash_output # 237μs -> 235μs (1.02% faster)
    # All elements should be equal (since einsum of ones)
    val = params["weight"].flatten()[0].item()

# Edge Test Cases

def test_zero_alpha():
    # Test with alpha=0, scale should be 0
    up = torch.ones(2, 2)
    down = torch.ones(2, 2)
    bias = torch.ones(2)
    orig_weight = torch.ones(2, 2)
    orig_bias = torch.ones(2)
    lora = FluxControlLoRALayer(up, None, down, 0.0, bias)
    codeflash_output = lora.get_parameters({"weight": orig_weight, "bias": orig_bias}, weight=1.0); params = codeflash_output # 45.2μs -> 36.4μs (24.1% faster)

def test_none_alpha():
    # Test with alpha=None, scale should be 1.0
    up = torch.ones(2, 2)
    down = torch.ones(2, 2)
    bias = torch.ones(2)
    orig_weight = torch.ones(2, 2)
    orig_bias = torch.ones(2)
    lora = FluxControlLoRALayer(up, None, down, None, bias)
    codeflash_output = lora.get_parameters({"weight": orig_weight, "bias": orig_bias}, weight=3.0); params = codeflash_output # 34.4μs -> 27.6μs (24.5% faster)
    # scale = 1.0
    expected_weight = torch.full((2, 2), 2.0) * 3.0
    expected_bias = torch.ones(2) * 3.0

def test_none_rank():
    # Test when _rank returns None (simulate by monkeypatching)
    class DummyLora(FluxControlLoRALayer):
        def _rank(self):
            return None
    up = torch.ones(2, 2)
    down = torch.ones(2, 2)
    bias = torch.ones(2)
    orig_weight = torch.ones(2, 2)
    orig_bias = torch.ones(2)
    lora = DummyLora(up, None, down, 5.0, bias)
    codeflash_output = lora.get_parameters({"weight": orig_weight, "bias": orig_bias}, weight=2.0); params = codeflash_output # 36.5μs -> 29.3μs (24.4% faster)
    # scale = 1.0
    expected_weight = torch.full((2, 2), 2.0) * 2.0
    expected_bias = torch.ones(2) * 2.0

def test_bias_none_in_orig_parameters():
    # Test when orig_parameters does not contain "bias"
    up = torch.ones(2, 2)
    down = torch.ones(2, 2)
    lora = FluxControlLoRALayer(up, None, down, 2.0, None)
    orig_weight = torch.ones(2, 2)
    codeflash_output = lora.get_parameters({"weight": orig_weight}, weight=1.0); params = codeflash_output # 32.1μs -> 23.5μs (36.4% faster)

def test_weight_zero():
    # Test with weight=0, output should be all zeros
    up = torch.ones(2, 2)
    down = torch.ones(2, 2)
    bias = torch.ones(2)
    orig_weight = torch.ones(2, 2)
    orig_bias = torch.ones(2)
    lora = FluxControlLoRALayer(up, None, down, 2.0, bias)
    codeflash_output = lora.get_parameters({"weight": orig_weight, "bias": orig_bias}, weight=0.0); params = codeflash_output # 35.2μs -> 27.3μs (28.9% faster)

def test_up_down_different_ranks():
    # up.shape[1] != down.shape[0], triggers fuse_weights
    up = torch.ones(2, 4)
    down = torch.ones(8, 2)
    orig_weight = torch.ones(2, 2)
    lora = FluxControlLoRALayer(up, None, down, 8.0, None)
    codeflash_output = lora.get_parameters({"weight": orig_weight}, weight=1.0); params = codeflash_output # 58.5μs -> 59.4μs (1.43% slower)

def test_mid_tensor_shape_mismatch():
    # Test with mid tensor not matching up/down shapes (should raise RuntimeError)
    up = torch.ones(2, 2)
    down = torch.ones(2, 2)
    mid = torch.ones(3, 2, 1, 1)  # mismatch on first dim
    orig_weight = torch.ones(2, 2, 1, 1)
    lora = FluxControlLoRALayer(up, mid, down, 2.0, None)
    with pytest.raises(RuntimeError):
        lora.get_parameters({"weight": orig_weight}, weight=1.0)

def test_missing_weight_in_orig_parameters():
    # Should raise KeyError if "weight" not in orig_parameters
    up = torch.ones(2, 2)
    down = torch.ones(2, 2)
    lora = FluxControlLoRALayer(up, None, down, 2.0, None)
    with pytest.raises(KeyError):
        lora.get_parameters({"bias": torch.ones(2)}, weight=1.0) # 3.33μs -> 3.09μs (7.78% faster)

def test_bias_shape_mismatch():
    # If bias is present, but shape mismatches, should still multiply
    up = torch.ones(2, 2)
    down = torch.ones(2, 2)
    bias = torch.ones(3)
    orig_weight = torch.ones(2, 2)
    orig_bias = torch.ones(3)
    lora = FluxControlLoRALayer(up, None, down, 2.0, bias)
    codeflash_output = lora.get_parameters({"weight": orig_weight, "bias": orig_bias}, weight=1.0); params = codeflash_output # 55.1μs -> 46.1μs (19.6% faster)

# Large Scale Test Cases

def test_large_up_down():
    # Test with large up/down matrices (but <100MB)
    up = torch.ones(32, 32)
    down = torch.ones(32, 32)
    orig_weight = torch.ones(32, 32)
    alpha = 32.0
    lora = FluxControlLoRALayer(up, None, down, alpha, None)
    codeflash_output = lora.get_parameters({"weight": orig_weight}, weight=1.0); params = codeflash_output # 34.9μs -> 27.5μs (26.7% faster)
    # scale = 1.0
    expected_weight = torch.full((32, 32), 32.0)

def test_large_mid_tensor():
    # Test with large mid tensor
    up = torch.ones(16, 16)
    down = torch.ones(16, 16)
    mid = torch.ones(16, 16, 1, 1)
    orig_weight = torch.ones(16, 16, 1, 1)
    alpha = 16.0
    lora = FluxControlLoRALayer(up, mid, down, alpha, None)
    codeflash_output = lora.get_parameters({"weight": orig_weight}, weight=1.0); params = codeflash_output # 243μs -> 239μs (1.46% faster)
    # All elements should be equal
    val = params["weight"].flatten()[0].item()

def test_large_bias():
    # Test with large bias
    up = torch.ones(32, 32)
    down = torch.ones(32, 32)
    bias = torch.ones(32)
    orig_weight = torch.ones(32, 32)
    orig_bias = torch.ones(32)
    alpha = 32.0
    lora = FluxControlLoRALayer(up, None, down, alpha, bias)
    codeflash_output = lora.get_parameters({"weight": orig_weight, "bias": orig_bias}, weight=2.0); params = codeflash_output # 43.8μs -> 37.7μs (16.1% faster)
    expected_bias = torch.ones(32) * 2.0

def test_large_scale_fuse_weights():
    # up.shape[1] != down.shape[0], triggers fuse_weights, large matrices
    up = torch.ones(32, 64)
    down = torch.ones(128, 32)
    orig_weight = torch.ones(32, 32)
    lora = FluxControlLoRALayer(up, None, down, 64.0, None)
    codeflash_output = lora.get_parameters({"weight": orig_weight}, weight=1.0); params = codeflash_output # 61.4μs -> 64.7μs (5.08% slower)

def test_large_weight_value():
    # Test with large weight value
    up = torch.ones(16, 16)
    down = torch.ones(16, 16)
    orig_weight = torch.ones(16, 16)
    alpha = 16.0
    lora = FluxControlLoRALayer(up, None, down, alpha, None)
    codeflash_output = lora.get_parameters({"weight": orig_weight}, weight=100.0); params = codeflash_output # 34.4μs -> 26.5μs (30.0% faster)
    expected_weight = torch.full((16, 16), 16.0) * 100.0
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-FluxControlLoRALayer.get_parameters-mhnaw21w and push.