import svetoviz_webgpu as sv

# 1. Define Layer Parameters
# p=0.5: 50% probability of an element being zeroed
dropout = nn.Dropout(p=0.5)

# 2. Prepare Input Tensor (Batch=1, Features=20)
input_tensor = torch.ones(1, 20)

def terminal_callback(buffer, message, images, files):
    # Dropout behavior is only active in .train() mode
    dropout.train()
    output = dropout(input_tensor)

    # Calculate scaling factor: 1 / (1 - p)
    scale_factor = 1 / (1 - dropout.p)

    buffer.send_system_message(f"Zeroed Elements: {(output == 0).sum().item()} / 20")
    buffer.send_system_message(f"Scaling Factor: {scale_factor}x")
    buffer.send_system_message("Logic: Remaining elements are scaled to preserve the expected sum.")

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