code: def solve_cubic(coeffs: Tensor) -> Tensor:
      r"""Solve given cubic equation.
  
      The function takes the coefficients of cubic equation and returns
      the real roots.
  
      .. math:: coeffs[0]x^3 + coeffs[1]x^2 + coeffs[2]x + coeffs[3] = 0
  
      Args:
          coeffs : The coefficients cubic equation : `(B, 4)`
  
      Returns:
          A tensor of shape `(B, 3)` containing the real roots to the cubic equation.
  
      Example:
          >>> coeffs = torch.tensor([[32., 3., -11., -6.]])
          >>> roots = solve_cubic(coeffs)
  
      .. note::
         In cases where a cubic polynomial has only one or two real roots, the output for the non-real
         roots should be represented as 0. Thus, the output for a single real root should be in the
         format [real_root, 0, 0], and for two real roots, it should be [real_root_1, real_root_2, 0].
      """
      KORNIA_CHECK_SHAPE(coeffs, ["B", "4"])
  
      _PI = torch.tensor(
          math.pi, device=coeffs.device, dtype=coeffs.dtype
      )
  
      # Coefficients of cubic equation
      a = coeffs[:, 0]  # coefficient of x^3
      b = coeffs[:, 1]  # coefficient of x^2
      c = coeffs[:, 2]  # coefficient of x
      d = coeffs[:, 3]  # constant term
  
      solutions = zeros(
          (len(coeffs), 3), device=a.device, dtype=a.dtype
      )
  
      mask_a_zero = a == 0
      mask_b_zero = b == 0
      mask_c_zero = c == 0
  
      # Zero order cases are automatically handled since solutions is initialized with zeros.
      # No need for explicit handling of mask_zero_order as solutions already contains zeros by default.
  
      mask_first_order = mask_a_zero & mask_b_zero & ~mask_c_zero
      mask_second_order = mask_a_zero & ~mask_b_zero & ~mask_c_zero
  
      if torch.any(mask_second_order):
          solutions[mask_second_order, 0:2] = solve_quadratic(
              coeffs[mask_second_order, 1:]
          )
  
      if torch.any(mask_first_order):
          solutions[mask_first_order, 0] = torch.tensor(
              1.0, device=a.device, dtype=a.dtype
          )
  
      # Normalized form x^3 + a2 * x^2 + a1 * x + a0 = 0
      inv_a = 1.0 / a[~mask_a_zero]
      b_a = inv_a * b[~mask_a_zero]
      b_a2 = b_a * b_a
  
      c_a = inv_a * c[~mask_a_zero]
      d_a = inv_a * d[~mask_a_zero]
  
      # Solve the cubic equation
      Q = (3 * c_a - b_a2) / 9
      R = (9 * b_a * c_a - 27 * d_a - 2 * b_a * b_a2) / 54
      Q3 = Q * Q * Q
      D = Q3 + R * R
      b_a_3 = (1.0 / 3.0) * b_a
  
      a_Q_zero = ones_like(a)
      a_R_zero = ones_like(a)
      a_D_zero = ones_like(a)
  
      a_Q_zero[~mask_a_zero] = Q
      a_R_zero[~mask_a_zero] = R
      a_D_zero[~mask_a_zero] = D
  
      # Q == 0
      mask_Q_zero = (Q == 0) & (R != 0)
      mask_Q_zero_solutions = (a_Q_zero == 0) & (a_R_zero != 0)
  
      if torch.any(mask_Q_zero):
          x0_Q_zero = (
              torch.pow(2 * R[mask_Q_zero], 1 / 3) - b_a_3[mask_Q_zero]
          )
          solutions[mask_Q_zero_solutions, 0] = x0_Q_zero
  
      mask_QR_zero = (Q == 0) & (R == 0)
      mask_QR_zero_solutions = (a_Q_zero == 0) & (a_R_zero == 0)
  
      if torch.any(mask_QR_zero):
          solutions[mask_QR_zero_solutions] = stack(
              [
                  -b_a_3[mask_QR_zero],
                  -b_a_3[mask_QR_zero],
                  -b_a_3[mask_QR_zero],
              ],
              dim=1,
          )
  
      # D <= 0
      mask_D_zero = (D <= 0) & (Q != 0)
      mask_D_zero_solutions = (a_D_zero <= 0) & (a_Q_zero != 0)
  
      if torch.any(mask_D_zero):
          theta_D_zero = torch.acos(
              R[mask_D_zero] / torch.sqrt(-Q3[mask_D_zero])
          )
          sqrt_Q_D_zero = torch.sqrt(-Q[mask_D_zero])
          x0_D_zero = (
              2 * sqrt_Q_D_zero * cos(theta_D_zero / 3.0)
              - b_a_3[mask_D_zero]
          )
          x1_D_zero = (
              2 * sqrt_Q_D_zero * cos((theta_D_zero + 2 * _PI) / 3.0)
              - b_a_3[mask_D_zero]
          )
          x2_D_zero = (
              2 * sqrt_Q_D_zero * cos((theta_D_zero + 4 * _PI) / 3.0)
              - b_a_3[mask_D_zero]
          )
          solutions[mask_D_zero_solutions] = stack(
              [x0_D_zero, x1_D_zero, x2_D_zero], dim=1
          )
  
      a_D_positive = zeros_like(a)
      a_D_positive[~mask_a_zero] = D
      # D > 0
      mask_D_positive_solution = (a_D_positive > 0) & (a_Q_zero != 0)
      mask_D_positive = (D > 0) & (Q != 0)
      if torch.any(mask_D_positive):
          AD = zeros_like(R)
          BD = zeros_like(R)
          R_abs = torch.abs(R)
          mask_R_positive = R_abs > 1e-16
          if torch.any(mask_R_positive):
              AD[mask_R_positive] = torch.pow(
                  R_abs[mask_R_positive]
                  + torch.sqrt(D[mask_R_positive]),
                  1 / 3,
              )
              mask_R_positive_ = R < 0
  
              if torch.any(mask_R_positive_):
                  AD[mask_R_positive_] = -AD[mask_R_positive_]
  
              BD[mask_R_positive] = (
                  -Q[mask_R_positive] / AD[mask_R_positive]
              )
          x0_D_positive = (
              AD[mask_D_positive]
              + BD[mask_D_positive]
              - b_a_3[mask_D_positive]
          )
          solutions[mask_D_positive_solution, 0] = x0_D_positive
  
      return solutions
  
file_path: kornia/kornia/geometry/solvers/polynomial_solver.py
imports: import math
  import torch
  from kornia.core import Tensor, cos, ones_like, stack, zeros, zeros_like
  from kornia.core.check import KORNIA_CHECK_SHAPE
label_desc: Use image node_id 2 for calling a global function with example usage: solve_cubic(coeffs) and returns: solutions