def adamw_stochastic_bf16_(  # noqa: C901
        params: torch.Tensor,
        grads: torch.Tensor,
        exp_avg: torch.Tensor,
        exp_avg_sq: torch.Tensor,
        lr: float,
        beta1: float,
        beta2: float,
        eps: float,
        weight_decay: float,
        step: int,
        generator: torch.Generator | None = None
) -> None:
    """
    Performs a single in-place AdamW optimization step.

    It is specifically designed for scenarios where parameters are stored in BFloat16.

    To mitigate precision loss during the parameter update, it utilizes stochastic rounding when casting
    FP32 calculation results back to BFloat16.

    This function supports mixed precision for gradients and optimizer states (they can be
    either FP32 or BFloat16).

    Args:
        params: The tensor of model parameters to update. Must be BFloat16 and contiguous.
        grads: The gradient tensor.
        exp_avg: The exponential moving average of gradient values (first moment).
        exp_avg_sq: The exponential moving average of squared gradient values (second moment).
        lr: The learning rate.
        beta1: Decay rate for the first moment estimate.
        beta2: Decay rate for the second moment estimate.
        eps: Term added to the denominator to improve numerical stability.
        weight_decay: Weight decay coefficient.
        step: The current optimization step count, used for bias correction.
        generator: PyTorch random number generator used to create the seed for stochastic rounding.

    Raises:
        ValueError: If main parameters are not BFloat16, if input tensor shapes do not match,
            if input tensors are not contiguous (for those that require in-place modification),
            if the optimizer states (exp_avg, exp_avg_sq) have different dtypes.
    """

    # check shape equality
    if grads.shape != params.shape:
        raise ValueError("Shape mismatch between grads and params.")

    if exp_avg.shape != params.shape:
        raise ValueError("Shape mismatch between exp_avg state and params.")

    if exp_avg_sq.shape != params.shape:
        raise ValueError("Shape mismatch between exp_avg_sq state and params.")

    # check params
    if params.dtype != torch.bfloat16:
        raise ValueError("Params must be BFloat16 for this kernel.")

    if not params.is_contiguous():
        raise ValueError("Params must be contiguous since it is an in-place kernel.")

    # check grads
    if not grads.is_contiguous():
        grads = grads.contiguous()

    # check states
    if not exp_avg.is_contiguous():
        raise ValueError("Exp_avg state must be contiguous since it is an in-place kernel.")

    if not exp_avg_sq.is_contiguous():
        raise ValueError("Exp_avg_sq state must be contiguous since it is an in-place kernel.")

    if exp_avg.dtype != exp_avg_sq.dtype:
        raise ValueError("States have different dtypes.")

    n_elements = params.numel()

    grad_is_bf16 = (grads.dtype == torch.bfloat16)
    state_is_bf16 = (exp_avg.dtype == torch.bfloat16)

    # Generate random seed
    seed = torch.randint(
        0, 2 ** 31 - 1, (1,),
        device="cpu",
        generator=generator
    ).item()

    def _grid(meta: dict[str, int]) -> tuple[int, ...]:
        return (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)

    _adamw_stochastic_bf16_kernel[_grid](
        params,
        grads,
        exp_avg,
        exp_avg_sq,

        n_elements,

        lr,
        beta1,
        beta2,
        eps,
        weight_decay,
        step,
        seed,

        GRAD_IS_BF16=grad_is_bf16,
        STATE_IS_BF16=state_is_bf16
    )

def copy_fp32_to_bf16_stochastic_(
        target: torch.Tensor,
        source: torch.Tensor,
        generator: torch.Generator | None = None
) -> torch.Tensor:
    """
    Copies elements from a Float32 tensor to a BFloat16 tensor using stochastic rounding.

    Unlike standard round-to-nearest casting, stochastic rounding probabilistically rounds
    numbers up or down based on the value of the bits being truncated. This preserves the
    expected value of the tensor (E[round(x)] = x), which is crucial for accumulating
    gradients or parameters in low precision without stagnation.

    This operation is performed in-place on the target tensor.

    Args:
        target: The output tensor where results are written. Must be of type BFloat16
            and contiguous.
        source: The input tensor containing values to copy. Must be of type Float32.
        generator: An optional PyTorch RNG generator to strictly control the random
            noise used for rounding.

    Returns:
        The target tensor, modified in-place.

    Raises:
        ValueError: If target is not contiguous, if source/target shapes do not match,
            or if dtypes are not FP32 and BF16 respectively.
    """

    if not source.is_contiguous():
        source = source.contiguous()

    if not target.is_contiguous():
        raise ValueError("Since this is an in-place operation, target should be a contiguous tensor!")

    if source.shape != target.shape:
        raise ValueError("Source and Target Tensors are of different shapes")

    if source.dtype != torch.float32:
        raise ValueError("Source must be Float32")
    if target.dtype != torch.bfloat16:
        raise ValueError("Target must be BFloat16")

    n_elements = source.numel()

    # Generate a random seed for this specific kernel launch
    seed = torch.randint(
        0, 2 ** 31 - 1, (1,),
        device="cpu",
        generator=generator
    ).item()

    def _grid(meta: dict[str, int]) -> tuple[int, ...]:
        return (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)

    _copy_fp32_to_bf16_kernel[_grid](
        source,
        target,
        n_elements,
        seed
    )
    return target