if (axis < 0) {
    axis += rank;
  }
  return b->getI64IntegerAttr(axis);
}

// Returns a PrecisionConfig as an array attribute based on whether TF32
// execution is enabled
static ArrayAttr GetPrecisionConfig(Builder *builder) {
  mlir::mhlo::Precision precision = tsl::tensor_float_32_execution_enabled()
                                        ? mhlo::Precision::DEFAULT