fixes

attn.py

@@ -0,0 +1,206 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+try:
+    from flash_attn import flash_attn_func
+except ImportError as e:
+    print(
+        f"Unable to import Triton-based flash attention: {e}. No alternative currently available."
+    )
+
+
+def nearest_power_of_two(x: int, round_up: bool = False) -> int:
+    return (
+        1 << math.floor(math.log2(x)) if not round_up else 1 << math.ceil(math.log2(x))
+    )
+
+def _generate_slopes(self, n: int):
+        start = 2 ** (-(2 ** -(math.log2(n) - 3)))
+        return [start * (start**i) for i in range(n)]
+
+def _get_alibi_slopes(self, n_heads: int, interpolation_factor: float = 0.25):
+    # If n_heads is a power of 2, generate slopes directly
+    if math.log2(n_heads).is_integer():
+        slopes = self._generate_slopes(n_heads)
+    else:
+        # Get slopes for the nearest power of two
+        n = nearest_power_of_two(n_heads, round_up=False)
+        slopes_power_of_two = self._generate_slopes(n)
+
+        # Generate extra slopes
+        extra_slopes = self._generate_slopes(2 * n)
+        extra_slopes_trunc = extra_slopes[0::2][: n_heads - n]
+        slopes = slopes_power_of_two + extra_slopes_trunc
+    slopes = torch.tensor(slopes, device=self.device)
+    slopes = slopes * interpolation_factor  # https://arxiv.org/pdf/2310.13017
+    return slopes
+
+
+def precompute_freqs_cis(head_dim: int, max_seq_len: int, theta: float = 10000.0):    
+    # For half the dimensions, build the scale factor:
+    freq_seq = torch.arange(0, head_dim, 2).float() / head_dim
+    freqs = 1.0 / (theta ** freq_seq)
+
+    # Outer product with positions
+    t = torch.arange(max_seq_len, dtype=torch.float32)
+    angles = torch.outer(t, freqs)
+    
+    # Build a complex exponential e^{i * theta}
+    freqs_cis = torch.polar(
+        torch.ones_like(angles),
+        angles
+    )
+    return freqs_cis
+
+
+def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
+    """
+    x is [B, n_heads, seq_len, head_dim_as_complex],
+    so we want to broadcast freqs_cis from [max_seq_len, half_dim]
+    to [1, 1, seq_len, half_dim].
+    """
+    seq_len = x.shape[2]
+    freqs_cis = freqs_cis[:seq_len]  # slice down to current seq_len
+    return freqs_cis.view(1, 1, seq_len, -1)
+
+
+def apply_rotary_emb(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    freqs_cis: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    # Convert real -> complex by grouping last dim in pairs
+    # shape => [B, n_heads, seq_len, head_dim//2, 2] => complex => [B, n_heads, seq_len, head_dim//2]
+    xq_complex = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
+    xk_complex = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
+
+    # Broadcast the frequencies to match [B, n_heads, seq_len, head_dim//2]
+    freqs_cis = reshape_for_broadcast(freqs_cis, xq_complex)
+
+    # Multiply => apply rotation
+    xq_complex = xq_complex * freqs_cis
+    xk_complex = xk_complex * freqs_cis
+
+    # Convert back to real => shape [B, n_heads, seq_len, head_dim]
+    xq_out = torch.view_as_real(xq_complex).reshape(*xq.shape)
+    xk_out = torch.view_as_real(xk_complex).reshape(*xk.shape)
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+
+class Attention(nn.Module):
+    def __init__(self, config):
+        super(Attention, self).__init__()
+        self.dim, self.num_heads = config.dim, config.num_heads
+        assert config.dim % config.num_heads == 0, f"dim ({self.dim}) must be divisible num_heads ({self.num_heads})"
+        self.head_dim = config.dim // config.num_heads
+
+        self.c_attn = nn.Linear(self.dim, 3*self.dim, bias=config.bias)
+        self.c_proj = nn.Linear(config.dim, config.dim, bias=config.bias)
+        self.c_proj.SCALE_INIT = 1
+
+        self.alibi_slopes = self._get_alibi_slopes(self.num_heads) if config.use_alibi else None
+        self.window_size = config.window_size
+        self.softcap = config.softcap
+
+        self.dropout = config.dropout
+        self.resid_dropout = nn.Dropout(self.dropout)
+
+    def _generate_slopes(self, n: int):
+            start = 2 ** (-(2 ** -(math.log2(n) - 3)))
+            return [start * (start**i) for i in range(n)]
+
+    def _get_alibi_slopes(self, num_heads: int, interpolation_factor: float = 0.25):
+        # If n_heads is a power of 2, generate slopes directly
+        if math.log2(num_heads).is_integer():
+            slopes = self._generate_slopes(num_heads)
+        else:
+            # Get slopes for the nearest power of two
+            n = nearest_power_of_two(num_heads, round_up=False)
+            slopes_power_of_two = self._generate_slopes(n)
+
+            # Generate extra slopes
+            extra_slopes = self._generate_slopes(2 * n)
+            extra_slopes_trunc = extra_slopes[0::2][: num_heads - n]
+            slopes = slopes_power_of_two + extra_slopes_trunc
+        slopes = torch.tensor(slopes, device=torch.device("cuda"))
+        slopes = slopes * interpolation_factor  # https://arxiv.org/pdf/2310.13017
+        return slopes
+
+    def forward(
+        self,
+        x: torch.Tensor = None,
+        q: torch.Tensor = None,
+        k: torch.Tensor = None,
+        v: torch.Tensor = None,
+        freqs_cis: torch.Tensor = None,
+    ) -> torch.Tensor:
+        if x is not None:
+            q = k = v = x
+        if any(t is None for t in [q, k, v]):
+            raise ValueError("Must provide either x for self-attention or q/k/v for cross-attention.")
+
+        bsz, q_len, dim = q.shape
+        _, k_len, _ = k.shape
+        _, v_len, _ = v.shape
+
+        qkv = self.c_attn(x)
+        q, k, v = torch.chunk(qkv, 3, dim=2)
+
+        q = q.view(bsz, q_len, self.num_heads, self.head_dim)
+        k = k.view(bsz, k_len, self.num_heads, self.head_dim)
+        v = v.view(bsz, v_len, self.num_heads, self.head_dim)
+
+        if self.alibi_slopes is None: # Use either ALiBi or RoPE
+            q, k = apply_rotary_emb(q, k, freqs_cis=freqs_cis)
+
+        y = flash_attn_func(  # https://arxiv.org/pdf/2307.08691
+            q=q, k=k, v=v,
+            dropout_p=self.dropout if self.training else 0.0,
+            causal=True,
+            window_size=(self.window_size, 0), # Set to config.seq_len if full attention
+            alibi_slopes=self.alibi_slopes, # https://arxiv.org/pdf/2108.12409
+            softcap=self.softcap,  # https://arxiv.org/pdf/2408.00118
+        )
+
+        y = y.contiguous().view(bsz, q_len, -1)
+        y = self.resid_dropout(self.c_proj(y))
+        return y
+
+
+class MLP(nn.Module):
+    def __init__(self, config):
+        # https://arxiv.org/pdf/2002.05202
+        super().__init__()
+        self.hidden_size = config.dim
+        self.intermediate_size = config.dim * config.mlp_scale
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.bias)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.bias)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.bias)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x):
+        gate = self.gate_proj(x)
+        gate = F.gelu(gate, approximate="tanh")
+        up = self.up_proj(x)
+        fuse = gate * up
+        outputs = self.down_proj(fuse)
+        outputs = self.dropout(outputs)
+        return outputs
+
+
+class AttentionLayer(nn.Module):
+    def __init__(self, config) -> None:
+        super(AttentionLayer, self).__init__()
+        self.attn_norm = nn.RMSNorm(config.dim)
+        self.attn = Attention(config=config)
+        self.mlp_norm = nn.RMSNorm(config.dim)
+        self.mlp = MLP(config)
+
+    def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor=None) -> torch.Tensor:
+        x = x + self.attn(x=self.attn_norm(x), freqs_cis=freqs_cis)
+        x = x + self.mlp(self.mlp_norm(x))
+        return x

causal_conv1d_compilable.py

@@ -211,5 +211,4 @@ if __name__ == "__main__":
 
     print(out.min(), out.max(), out.mean(), out.std())
     print(x.grad.min(), x.grad.max(), x.grad.mean(), x.grad.std())
-    print(weight.grad.min(), weight.grad.max(), weight.grad.mean(), weight.grad.std())
-
+    print(weight.grad.min(), weight.grad.max(), weight.grad.mean(), weight.grad.std())

model.py

@@ -1,3 +1,12 @@
+"""
+
+Adapted from Meta's Lingua repository:
+    - https://github.com/facebookresearch/lingua/blob/main/apps/mamba/core_mamba.py
+    - https://github.com/facebookresearch/lingua/blob/main/apps/mamba/mamba.py
+
+"""
+
+import json
 import  math
 
 import torch
@@ -6,13 +15,19 @@ import torch.nn.functional as F
 
 from enum import Enum
 from dataclasses import dataclass, field
+
 from causal_conv1d.causal_conv1d_varlen import causal_conv1d_varlen_states
+from .causal_conv1d_compilable import (
+    causal_conv1d_fn as causal_conv1d_fn, 
+    causal_conv1d_update as causal_conv1d_update
+)
+
 from mamba_ssm.ops.triton.selective_state_update import selective_state_update
 from mamba_ssm.ops.triton.ssd_combined import mamba_split_conv1d_scan_combined
 
-from .causal_conv1d_compilable import causal_conv1d_fn, causal_conv1d_update
 from .ssm_compilable import mamba_chunk_scan_combined
 from .norms import build_norm
+from .attn import AttentionLayer
 
 
 class InitStdFactor(Enum):
@@ -154,9 +169,7 @@ class SSM(nn.Module):
         if self.learnable_init_states:
             self.init_states = nn.Parameter(torch.zeros(self.num_heads, self.head_dim, self.state_dim))
 
-        # Can also just use nn.RMSNorm
         self.norm = build_norm(config.norm_type, dim=self.hidden_dim, eps=self.norm_eps)
-        
         self.output = nn.Linear(self.hidden_dim, self.dim, bias=self.bias)
 
     def _causal_conv(
@@ -320,7 +333,7 @@ class SSM(nn.Module):
             ),
             dt_softplus=True,
         ).unsqueeze(0)
-        
+
         return y
 
     def forward(
@@ -502,8 +515,16 @@ class BaseMamba(nn.Module):
         self.init_std_factor = InitStdFactor(config.init_std_factor)
 
         self.layers = nn.ModuleList()
-        for _ in range(config.num_layers):
-            self.layers.append(MambaBlock(config))
+        for layer_idx in range(config.num_layers):
+            # For more complex %-split arrangements, see https://arxiv.org/pdf/2406.07887
+            if layer_idx % 2 == 0:
+                self.layers.append(MambaBlock(config))
+            else:
+                self.layers.append(
+                    AttentionLayer(config)
+                    if config.use_attn
+                    else (MambaBlock(config))
+                )
 
     def forward(
         self,
@@ -536,6 +557,7 @@ class BaseMamba(nn.Module):
 
 @dataclass
 class Mamba2Config(BaseMambaConfig):
+    bsz: int = 2
     seed: int = 1337
 
     vocab_size: int = -1 # Will error if unchanged, makes you double check!
@@ -573,10 +595,10 @@ class Mamba2(BaseMamba):
 
     def _get_num_params(self):
         n_params = sum(p.numel() for p in self.parameters())
+
         if hasattr(self, "pos_emb") and self.pos_emb is not None:
             n_params -= self.pos_emb.weight.numel()
-        if self.tok_emb.weight is not self.output.weight:
-            n_params -= self.tok_emb.weight.numel()
+
         return n_params
 
     def forward(
@@ -657,192 +679,108 @@ class Mamba2(BaseMamba):
         return cls(config)
 
 
-def get_mamba2_flops(
-    seq_len: int,
-    dim: int,
-    num_layers: int,
-    vocab_size: int,
-    ffn_multiplier: float = 2.0,
-    state_dim: int = 128,
-    conv_size: int = 4,
-    num_heads: int = 8,
-    num_groups: int = 1,
-    multiple_of: int = 256,
-    include_input_embedding: bool = True,
-    include_output_logits: bool = True,
-    forward_backward_multiplier: float = 1.0,
-) -> int:
-    """
-    Estimate the FLOPs for a Mamba-2 style model using a "Chinchilla-like" shape-based approach.
-
-    By default, this returns the forward-pass cost. If you want a rough
-    forward+backward estimate, set `forward_backward_multiplier=3.0` (common
-    rule-of-thumb for these models).
-
-    What gets counted:
-    • Hidden dimension is rounded up to 'multiple_of' = 256 (as in Mamba).
-    • Per-layer:
-        1) Input Linear: [dim → 2*hidden_dim + 2*(groups*state_dim) + num_heads]
-        2) Depthwise Conv1D: 2*(conv_dim * conv_size), where conv_dim=hidden_dim + 2*groups*state_dim
-        3) SSM selective scan: ~9*(dim*state_dim) (from Mamba dev discussion)
-        4) Output Linear: [hidden_dim → dim]
-    • Each layer’s cost is multiplied by (seq_len * num_layers).
-    • Optionally adds:
-        - The cost of the input embedding (treating it as a matmul: seq_len×vocab_size × vocab_size×dim).
-        - The cost of the final projection [dim → vocab_size].
-    • Finally scaled by `forward_backward_multiplier` if desired.
-
-    Args:
-        seq_len (int): Sequence length (number of tokens).
-        dim (int): Model (embedding) dimension.
-        num_layers (int): Number of Mamba layers.
-        vocab_size (int): Vocabulary size for final logits projection.
-        ffn_multiplier (float): FFN expansion ratio, e.g. 2.0 => hidden_dim=2×dim (rounded up).
-        state_dim (int): SSM state dimension (commonly 128).
-        conv_size (int): Kernel size for the depthwise conv1d (default=4).
-        num_heads (int): Number of heads (slightly affects input-lin out_dim).
-        num_groups (int): For "grouped" states in some Mamba variants (usually 1).
-        multiple_of (int): Round hidden_dim up to this multiple (commonly 256).
-        include_input_embedding (bool): If True, count the cost of an “embedding matmul”
-                                        for the input tokens => shape-based approach.
-        include_output_logits (bool): If True, count the cost of final [dim → vocab_size].
-        forward_backward_multiplier (float): E.g. 1.0 for forward only, 2.0 or 3.0 for forward+backward.
-
-    Returns:
-        int: Approximate total FLOPs (multiply-adds) for the selected pass(es),
-            as an integer.
-    """
-    # 0) Input embedding (optional)
-    flops_embedding = 0
-    if include_input_embedding:
-        flops_embedding = 2 * (seq_len * vocab_size * dim)
-
-    # 1) Round up hidden_dim
-    raw_hidden_dim = int(ffn_multiplier * dim)
-    hidden_dim = multiple_of * ((raw_hidden_dim + multiple_of - 1) // multiple_of)
-
-    # 2) Per-layer forward cost
-    out_dim_input = 2*hidden_dim + 2*(num_groups*state_dim) + num_heads
-    flops_input_linear = 2 * (dim * out_dim_input)
-    conv_dim = hidden_dim + 2*(num_groups*state_dim)
-    flops_conv = 2 * (conv_dim * conv_size)
-    flops_ssm = 9 * state_dim * dim
-    flops_output_linear = 2 * (hidden_dim * dim)
-    flops_layer = (flops_input_linear + flops_conv + flops_ssm + flops_output_linear)
-
-    # Multiply by #layers and sequence length
-    flops_layers = flops_layer * num_layers * seq_len
-
-    # 3) Final projection [dim → vocab_size] (optional)
-    flops_vocab = 0
-    if include_output_logits:
-        flops_vocab = 2 * (seq_len * dim * vocab_size)
-
-    # 4) Total forward FLOPs
-    flops_forward = flops_embedding + flops_layers + flops_vocab
-
-    # 5) Scale for forward+backward if desired
-    return int(flops_forward * forward_backward_multiplier)
-
-def get_mamba2_flops_per_token(
-    **kwargs
-) -> float:
-    """
-    Estimate FLOPs per token for a Mamba-2 style model.
-
-    This function extracts necessary parameters from kwargs and calculates the FLOPs per token.
-    
-    Args:
-        **kwargs: Dictionary containing model configuration parameters.
-    
-    Returns:
-        float: Approximate FLOPs per token.
-    """
-    defaults = {
-        'ffn_dim_multiplier': 2.0,
-        'state_dim': 128,
-        'conv_size': 4,
-        'num_heads': 8,
-        'num_groups': 1,
-        'multiple_of': 256,
-        'include_input_embedding': True,
-        'include_output_logits': True,
-        'forward_backward_multiplier': 1.0,
-    }
-    # Merge defaults
-    for k, v in defaults.items():
-        kwargs.setdefault(k, v)
-    # Mandatory keys
-    for required in ['seq_len', 'dim', 'num_layers', 'vocab_size']:
-        if required not in kwargs:
-            raise ValueError(f"Missing required parameter: {required}")
-
-    total_flops = get_mamba2_flops(
-        seq_len=kwargs['seq_len'],
-        dim=kwargs['dim'],
-        num_layers=kwargs['num_layers'],
-        vocab_size=kwargs['vocab_size'],
-        ffn_multiplier=kwargs['ffn_dim_multiplier'],
-        state_dim=kwargs['state_dim'],
-        conv_size=kwargs['conv_size'],
-        num_heads=kwargs['num_heads'],
-        num_groups=kwargs['num_groups'],
-        multiple_of=kwargs['multiple_of'],
-        include_input_embedding=kwargs['include_input_embedding'],
-        include_output_logits=kwargs['include_output_logits'],
-        forward_backward_multiplier=kwargs['forward_backward_multiplier'],
+# def main():
+#     x = torch.randn(2, 64, 192).cuda()
+
+#     config = Mamba2Config(vocab_size=200064, use_mem_eff_path=True)
+#     model2 = Mamba2(
+#         config=config,
+#     ).cuda()
+
+#     x_indices = torch.randint(0, config.vocab_size, (2, 64), dtype=torch.long).cuda()
+
+#     y2 = model2(x_indices)
+#     print("Mamba output: ", y2)
+#     print("Mean of Mamba output: ", y2.mean().item())
+#     print("Stddev of Mamba output: ", y2.std().item())
+
+# if __name__ == "__main__":
+#     main()
+
+if __name__ == '__main__':
+    x = torch.randn(2, 64, 192).cuda() # Removing this produces NaNs lol
+
+    config_path = "/scratch/gpfs/mn4560/hazan-lab/tensorized_filters/tensorized_filters/models/mamba/config.json"
+
+    with open(config_path, "r") as f:
+        config_data = json.load(f)
+
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+    elif torch.backends.mps.is_available():
+        device = torch.device("mps")
+    else:
+        device = torch.device("cpu")
+    print("Device:", device)
+
+    torch_dtype = getattr(torch, config_data["torch_dtype"])
+    print("Torch dtype:", torch_dtype)
+
+    dim = config_data["dim"]
+    num_heads = config_data["num_heads"]
+    num_layers = config_data["num_layers"]
+    vocab_size = config_data["vocab_size"]
+    bias = config_data["bias"]
+    state_dim = config_data["state_dim"]
+    num_groups = config_data["num_groups"]
+    conv_size = config_data.get("conv_size")
+    use_mem_eff_path = config_data.get("use_mem_eff_path")
+    dt_bias = config_data["dt_bias"]
+    D_has_head_dim = config_data["D_has_head_dim"]
+    learnable_init_states = config_data["learnable_init_states"]
+    ssm_chunk_size = config_data["ssm_chunk_size"]
+    weight_tying = config_data["weight_tying"]
+    ffn_dim_multiplier = config_data.get("ffn_dim_multiplier")
+    multiple_of = config_data["multiple_of"]
+    norm_eps = config_data["norm_eps"]
+    init_use_depth = config_data["init_use_depth"]
+    init_base_std = config_data.get("init_base_std")
+    init_std_factor = config_data["init_std_factor"]
+    use_attn = config_data["use_attn"]
+    softcap = config_data["softcap"]
+    torch_compile = config_data["torch_compile"]
+
+    configs = Mamba2Config(
+        dim=dim,
+        num_layers=num_layers,
+        num_heads=num_heads,
+        vocab_size=vocab_size,
+        bias=bias,
+        torch_dtype=torch_dtype,
+        state_dim=state_dim,
+        num_groups=num_groups,
+        conv_size=conv_size,
+        use_mem_eff_path=use_mem_eff_path,
+        dt_bias=dt_bias,
+        D_has_head_dim=D_has_head_dim,
+        learnable_init_states=learnable_init_states,
+        ssm_chunk_size=ssm_chunk_size,
+        weight_tying=weight_tying,
+        ffn_dim_multiplier=ffn_dim_multiplier,
+        multiple_of=multiple_of,
+        norm_eps=norm_eps,
+        init_use_depth=init_use_depth,
+        init_base_std=init_base_std,
+        init_std_factor=init_std_factor,
+        use_attn=use_attn,
+        softcap=softcap,
     )
-    flops_per_token = total_flops / kwargs['seq_len']
-
-    return flops_per_token
-
-
-# Optional policy for activation checkpointing. With None, we stick to the default (defined distributed.py: default_no_recompute_ops)
-def get_no_recompute_ops():
-    return {
-        torch.ops.aten.mm.default,
-        torch.ops.aten._scaled_mm.default,
-        torch.ops.c10d_functional.reduce_scatter_tensor.default,
-        torch.ops.mamba_ssm.ssm_chunk_scan_combined_fwd.default,
-
-        # For low-precision training, it's useful to always save the result of max(abs(tensor))
-        torch.ops.aten.abs.default,
-        torch.ops.aten.max.default,
-    }
-
-
-def main():
-    from mamba_ssm import Mamba2 as MambaRef
-
-    x = torch.randn(2, 64, 192).cuda()
-
-    # Create and run the first model
-    model = MambaRef(
-        d_model=192,
-        expand=2,
-        d_conv=4,
-        d_state=64,
-        headdim=48,
-    ).cuda()
-    y = model(x)
-    print("Mamba reference output: ", y)
-    print("Mean of MambaRef output: ", y.mean().item())
-    print("Stddev of MambaRef output: ", y.std().item())
 
-    # Create and run the second model
-    config = Mamba2Config(vocab_size=200064, use_mem_eff_path=True)
-    model2 = Mamba2(
-        config=config,
-    ).cuda()
+    print("Configs:")
+    for key, value in vars(configs).items():
+        print(f"  {key}: {value}")
 
-    # Fix: Convert x to torch.LongTensor
-    x_indices = torch.randint(0, config.vocab_size, (2, 64), dtype=torch.long).cuda()
+    model = Mamba2(configs).to(device=device)
 
-    y2 = model2(x_indices)
-    print("Mamba output: ", y2)
-    print("Mean of Mamba output: ", y2.mean().item())
-    print("Stddev of Mamba output: ", y2.std().item())
+    x = torch.randint(
+        0, configs.vocab_size, 
+        (config_data["bsz"], config_data["seq_len"]), 
+        dtype=torch.long
+    ).to(device)
 
-if __name__ == "__main__":
-    main()
+    outputs = model(x)
 
+    print("Output shape:", outputs.shape)
+    print("Sample output:", outputs[0, 0, :10])
+    print("Mean of Mamba output: ", outputs.mean().item())
+    print("Stddev of Mamba output: ", outputs.std().item())

norms.py

@@ -354,4 +354,4 @@ def fused_rms_norm_fn(
         x,
         weight,
         eps,
-    )
+    )

ssm_compilable.py

@@ -218,5 +218,4 @@ if __name__ == "__main__":
 
     out_ref = mamba_chunk_scan_combined_ref(x, dt, A, B, C, chunk_size, D=D, z=z, dt_bias=dt_bias)
 
-    print(out_ref.min(), out_ref.max(), out_ref.mean(), out_ref.std())
-
+    print(out_ref.min(), out_ref.max(), out_ref.mean(), out_ref.std())