import svetoviz_webgpu as sv

# 1. Define Layer Parameters
# in_features=128: The size of each input sample
# out_features=64: The size of each output sample
linear_layer = nn.Linear(in_features=128, out_features=64, bias=True)

# 2. Prepare Input Tensor (Batch=1, Features=128)
input_tensor = torch.randn(1, 128)

def terminal_callback(buffer, message, images, files):
    # Forward pass
    output = linear_layer(input_tensor)

    # Access weights and bias shapes
    weights_shape = list(linear_layer.weight.shape)
    bias_shape = list(linear_layer.bias.shape) if linear_layer.bias is not None else "None"

    buffer.send_system_message(f"Input Shape: {list(input_tensor.shape)}")
    buffer.send_system_message(f"Weight Matrix: {weights_shape} (Out x In)")
    buffer.send_system_message(f"Output Shape: {list(output.shape)}")

    # Calculate total parameters: (In * Out) + Out
    total_params = linear_layer.weight.numel() + (linear_layer.bias.numel() if linear_layer.bias is not None else 0)
    buffer.send_system_message(f"Total Learnable Parameters: {total_params}")

# 3. Start the interactive session
sv.pytorch_web(module=linear_layer, terminal_callback=terminal_callback)