IntegerAttr intAttr = mlir::dyn_cast_or_null<IntegerAttr>(attr);
  if (auto elementAttr = mlir::dyn_cast_or_null<ElementsAttr>(attr)) {
    SmallVector<uint64_t, 1> index(elementAttr.getShapedType().getRank(), 0);
    intAttr = elementAttr.getValues<IntegerAttr>()[index];
  }

  assert(intAttr && "Invalid attribute passed to GetHLOAxisFromTFAxis");

  int64_t axis = intAttr.getInt();
  if (axis < 0) {