auto slice_begin_const =
        rewriter->create<TFL::ConstOp>(split_op->getLoc(), slice_begin_attr);
    auto slice_size_const =
        rewriter->create<TFL::ConstOp>(split_op->getLoc(), slice_size_attr);

    auto slice_op = rewriter->create<TFL::SliceOp>(
        split_op->getLoc(), current_output_type, input, slice_begin_const,
        slice_size_const);

    // Rewire output.

  auto erf_op = rewriter.create<chlo::ErfOp>(mul_op.getLoc(), mul_op);
  auto add_op =
      rewriter.create<stablehlo::AddOp>(erf_op.getLoc(), erf_op, one_cst);
  auto lhs_mul_op = rewriter.create<stablehlo::MulOp>(
      half_cst.getLoc(), new_func.getArguments()[0], half_cst);
  auto output_mul_op = rewriter.create<stablehlo::MulOp>(lhs_mul_op.getLoc(),

    Value start_op = rewriter.create<TF::ConstOp>(iota_op.getLoc(), start);
    Value limit_op = rewriter.create<TF::ConstOp>(iota_op.getLoc(), limit);
    Value delta_op = rewriter.create<TF::ConstOp>(iota_op.getLoc(), delta);
    Value result = rewriter.create<TF::RangeOp>(iota_op.getLoc(), range_type,
                                                start_op, limit_op, delta_op);

      rewriter.create<func::FuncOp>(input_op->getLoc(), func_name, func_type);

  OpBuilder::InsertionGuard guard = OpBuilder::InsertionGuard(rewriter);
  ArrayRef<Type> inputs = quantization_func.getFunctionType().getInputs();
  Block* block = rewriter.createBlock(
      &quantization_func.getBody(), quantization_func.begin(), inputs,
      SmallVector<Location>(inputs.size(), quantization_func.getLoc()));
  func_input_arg = block->getArgument(0);

    auto min_scaled =
        rewriter.create<DivOp>(op.getLoc(), float_min, quant_to_float);
    auto min_scaled_sub =
        rewriter.create<SubOp>(op.getLoc(), quant_min, min_scaled);

    auto mid_rounded =
        rewriter.create<RoundOp>(op.getLoc(), scalar_ty, min_scaled_sub);

    auto nudged_zero_point_val = rewriter.create<ClipByValueOp>(
        op.getLoc(), scalar_ty, mid_rounded, quant_min, quant_max);

    auto range = rewriter.create<mhlo::SubtractOp>(op.getLoc(), limit, start);
    auto abs = rewriter.create<mhlo::AbsOp>(op.getLoc(), range);

    // Delta is not necessarily the same type as start and limit.
    auto abs_cast =
        rewriter.create<mhlo::ConvertOp>(op.getLoc(), compute_type, abs);
    auto delta_cast =
        rewriter.create<mhlo::ConvertOp>(op.getLoc(), compute_type, delta);

      auto tf_reduce_op = rewriter.create<BooleanReduce>(
          reduce_op.getLoc(), reduce_op->getResult(0).getType(), operand,
          reduction_indices,
          /*keep_dim=*/rewriter.getBoolAttr(false));
      auto tf_argreduce_op = rewriter.create<ArgReduce>(
          reduce_op.getLoc(), reduce_op->getResult(1).getType(), operand,
          reduction_indices);

        filter_constant_op->getLoc(), *rewriter.getContext(),
        filter_quantized_type.getScale(), filter_quantized_type.getZeroPoint(),
        /*narrow_range=*/true);
  }

  // Required because the quantized dimension is changed from 3 -> 0.
  auto new_filter_result_type = RankedTensorType::getChecked(
      filter_constant_op.getLoc(), /*shape=*/new_filter_shape,

              builder.create<ConstOp>(op->getLoc(), output_type, attribute);
        } else {
          attr_cst = builder.create<mlir::arith::ConstantOp>(
              op->getLoc(), cast<TypedAttr>(attribute));
        }
        new_operands.push_back(attr_cst);
      }
    }

    // Create the TFR call op
    auto new_op = builder.create<CallOp>(
        op->getLoc(), compose_func_type.getResults(),

  builder.insert(var_handle_op_in_init);
  auto const_op = builder.create<mlir::arith::ConstantOp>(
      session_init_func.getLoc(), tensor_attr.getType(), tensor_attr);

  builder.create<TF::AssignVariableOp>(
      session_init_func.getLoc(), llvm::ArrayRef<mlir::Type>{},
      llvm::ArrayRef<mlir::Value>{var_handle_op_in_init->getResult(0),
                                  const_op.getResult()});
}