if (val.getType() == t) return val;
        return rewriter.create<tensor::CastOp>(op.getLoc(), t, val);
      };
      last_val[0] = maybe_cast(const_val, op.getResult(3).getType());
      last_val[1] = maybe_cast(const_val, op.getResult(4).getType());
    }
    rewriter.replaceOp(
        op, {/*x_backprop=*/x_backprop,
             /*scale_backprop=*/scale_backprop,