Value offset = op.getShift();
    auto axis_size = b.create<mhlo::ConstantOp>(b.getIntegerAttr(
        getElementTypeOrSelf(offset.getType()), input_shape[axis]));
    offset = b.create<RemOp>(
        b.create<AddOp>(b.create<RemOp>(offset, axis_size), axis_size),
        axis_size);

    // Stack two copies of the dimension, then slice from the calculated
    // offset. This also works if shift is not constant.