stablehlo.cbrt crashes with complex numbers when StableHLO spec supports it

[mofeing]

from EnzymeAD/Reactant.jl#291
CC @wsmoses

StableHLO spec for cbrt specifies that it should support complex numbers but when I try to run it, it crashes with the following error:

terminate called after throwing an instance of 'xla::XlaRuntimeError'
  what():  UNKNOWN: <unknown>:0: error: Expected element type in shape to be floating for cbrt operation; got C128.: 
<unknown>:0: note: see current operation: "func.return"(%0) : (tensor<4xcomplex<f64>>) -> ()

MWE

The MLIR code for the MWE is the following:

module {
  func.func @main(%arg0: tensor<4xcomplex<f64>>) -> (tensor<4xcomplex<f64>>, tensor<4xcomplex<f64>>) {
    %0 = stablehlo.transpose %arg0, dims = [0] : (tensor<4xcomplex<f64>>) -> tensor<4xcomplex<f64>>
    %1 = stablehlo.cbrt %0 : tensor<4xcomplex<f64>>
    %2 = stablehlo.transpose %1, dims = [0] : (tensor<4xcomplex<f64>>) -> tensor<4xcomplex<f64>>
    %3 = stablehlo.transpose %0, dims = [0] : (tensor<4xcomplex<f64>>) -> tensor<4xcomplex<f64>>
    return %2, %3 : tensor<4xcomplex<f64>>, tensor<4xcomplex<f64>>
  }
}

You can forget about the transposes since we use them to accommodate Julia's column-major layout to MLIR's row-major layout and are later optimized away.

[akuegel]

Stablehlo is promising something that is (currently) not supported by XLA:

xla/xla/service/shape_inference.cc

Line 351 in 74da9de

case HloOpcode::kCbrt: // Complex cbrt is not implemented in either of the

@GleasonK do you know why stablehlo allows complex type? Do we have some wrong documentation about what types are supported somewhere in XLA?
The fix to shape inference (i.e. check for what is actually supported) was done in this PR: #18438

[majnemer]

I hesitated adding a complex variant of cbrt back in the day because the principal and real valued roots disagreed when the input was real. That said, we should just spec it as the principal root:

  r = abs(z);
  theta = atan2(imag(z), real(z));
  root_r = cbrt(r);
  root_theta = theta / 3.0;

  return root_r * (complex(cos(root_theta), sin(root_theta)));