CHECK-NEXT: %[[cst_4:.*]] = arith.constant 0 : index
      CHECK-NEXT: %[[elt:.*]] = tfr.get_element %[[Split]][%idx] : (!tfr.tensor_list, index) -> !tfr.tensor
      CHECK-NEXT: %[[cst_5:.*]] = arith.constant 1 : index
      CHECK-NEXT: %[[elt_1:.*]] = tfr.get_element %[[Split]][%idx_1] : (!tfr.tensor_list, index) -> !tfr.tensor
      CHECK-NEXT: constant true

  %v1 = tfr.get_element %values[%index] : (!tfr.tensor_list, index) -> !tfr.tensor
  %res = scf.if %eq -> !tfr.tensor {
    scf.yield %v1 : !tfr.tensor
  } else {
    %step = arith.index_cast %cst : i64 to index
    %end = arith.index_cast %n : i64 to index
    %reduce = scf.for %i = %step to %end step %step iter_args(%reduce_iter=%v1) -> !tfr.tensor {

                      %axis: i32 {tfr.name="axis"}) -> !tfr.tensor {
  %index = arith.constant 0 : index
  %cst = arith.constant 1 : i32
  %eq = arith.cmpi eq, %n, %cst : i32
  %v1 = tfr.get_element %values[%index] : (!tfr.tensor_list, index) -> !tfr.tensor
  %temp = tfr.call @tf__expand_dims(%v1, %axis) : (!tfr.tensor, i32) -> !tfr.tensor
  %res = scf.if %eq -> !tfr.tensor {
    scf.yield %temp : !tfr.tensor
  } else {

  %v1 = tfr.get_element %values[%index] : (!tfr.tensor_list, index) -> !tfr.tensor
  %res = scf.if %eq -> !tfr.tensor {
    scf.yield %v1 : !tfr.tensor
  } else {
    %step = arith.index_cast %cst : i64 to index
    %end = arith.index_cast %n : i64 to index
    %reduce = scf.for %i = %step to %end step %step iter_args(%reduce_iter=%v1) -> !tfr.tensor {

                      %axis: i32 {tfr.name="axis"}) -> !tfr.tensor {
  %index = arith.constant 0 : index
  %cst = arith.constant 1 : i32
  %eq = arith.cmpi eq, %n, %cst : i32
  %v1 = tfr.get_element %values[%index] : (!tfr.tensor_list, index) -> !tfr.tensor
  %temp = tfr.call @tf__my_expand_dims(%v1, %axis) : (!tfr.tensor, i32) -> !tfr.tensor
  %res = scf.if %eq -> !tfr.tensor {
    scf.yield %temp : !tfr.tensor

  %0 = "tfr.constant_tensor"(%arg0) : (i32) -> tensor<1x1xi32>
  func.return %0 : tensor<1x1xi32>
}

// -----

// CHECK-LABEL: get_element
func.func @get_element(%arg0: !tfr.tensor_list<T>) -> !tfr.tensor {
  %cst = "arith.constant"() {value = 1 : index} : () -> index
  %0 = tfr.get_element %arg0[%cst] : (!tfr.tensor_list<T>, index) -> !tfr.tensor
  func.return %0 : !tfr.tensor
}

// -----

// CHECK-LABEL: build_list

  let hasCanonicalizer = 1;

  let hasVerifier = 1;
}

def TFR_GetElementOp : TFR_Op<"get_element", [Pure]> {
  let description = [{
    The `get_element` operation extracts one tfr.tensor element from a
    tfr.tensor_list.

    Example:

    ```mlir
    %2 = tfr.get_element %1[%0] : (tfr.tensor, index) -> tfr.tensor
    ```
  }];

  let arguments = (ins

  %0 = "tfr.cast"(%arg0) : (tensor<3xf32>) -> !tfr.tensor
  %n = arith.constant 3: i32
  %split = tfr.call @tf__risc_split(%0, %n) : (!tfr.tensor, i32) -> !tfr.tensor_list
  %i0 = arith.constant 0: index
  %s0 = tfr.get_element %split[%i0] : (!tfr.tensor_list, index) -> !tfr.tensor
  %4 = "tfr.cast"(%s0) : (!tfr.tensor) -> tensor<f32>
  func.return %4 : tensor<f32>

    %arg0: !tfr.tensor,
    %arg1: !tfr.tensor_list,
    %arg2: index {tfr.name="A",tfr.default=1:index})
  -> !tfr.tensor_list {
  %0 = tfr.call @tf__intermediate(%arg0) : (!tfr.tensor) -> !tfr.tensor
  %1 = tfr.get_element %arg1[%arg2] : (!tfr.tensor_list, index) -> !tfr.tensor
  %2 = tfr.call @tf__intermediate(%1) : (!tfr.tensor) -> !tfr.tensor
  %3 = "tfr.build_list"(%0, %2) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor_list

        new_operands);

    // Replace the use of the old op. This is mapping the results from the
    // target TF ops to the TFR function returns. If the TFR function return is
    // a tensor_list, "tfr.get_element" op is used to extract the required TF
    // op result.
    llvm::SmallVector<Value, 4> new_results;
    for (auto res : llvm::enumerate(compose_func_type.getResults())) {