push.replaceAllUsesWith(push.getElem());
  OpBuilder builder(push);
  // Read the current buffer and size.
  auto stack_val =
      cutil::ReadLocalVariable(push.getHandle(), builder, push.getLoc());
  auto index =
      cutil::ReadLocalVariable(it->getSecond(), builder, push.getLoc());
  stack_val = cutil::SetElement(index, stack_val, push.getElem(), builder,
                                push.getLoc());

  auto local_var = read.getHandle();
  if (stats.count(local_var) == 0) {
    return read.emitOpError("unknown tensor array");
  }
  OpBuilder builder(read);
  auto buffer = cutil::ReadLocalVariable(local_var, builder, read.getLoc());
  auto index_reshape =
      cutil::ReshapeScalarToSizeType(builder, read.getIndex(), read.getLoc());
  auto elem = cutil::GetElement(index_reshape, buffer, builder, read.getLoc());

std::optional<RankedTensorType> GetElementTypeFromAccess(
    Value collection, ModuleOp module,
    llvm::function_ref<std::optional<Type>(Operation*)> infer_from_op);

// Creates a ReadVariableOp on a local variable.
Value ReadLocalVariable(Value local_var, OpBuilder builder, Location loc);

// Creates an AssignVariableOp on a local variable.
TF::AssignVariableOp WriteLocalVariable(Value local_var, Value value,

      if (elem_type && elem_type.hasStaticShape()) return elem_type;
    }
  }
  return std::nullopt;
}

// Creates a ReadVariableOp on a local variable.
Value ReadLocalVariable(Value local_var, OpBuilder builder, Location loc) {
  return builder
      .create<TF::ReadVariableOp>(
          loc,
          ArrayRef<Type>{mlir::cast<TF::ResourceType>(