diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
new file mode 100644
index 000000000..850db49a9
--- /dev/null
+++ b/.github/workflows/ci.yml
@@ -0,0 +1,217 @@
+name: Run all checks for pull requests
+
+on:
+  pull_request:
+    branches:
+      - main
+  workflow_dispatch:
+    inputs:
+      ref:
+        description: 'The git ref to build the package for'
+        required: false
+        default: ''
+        type: string
+
+# Precompute the ref if the workflow was triggered by a workflow dispatch rather than copying this logic repeatedly
+env:
+  ref: ${{ github.event_name == 'workflow_dispatch' && inputs.ref || null }}
+
+jobs:
+  eval: 
+    name: Evaluate changes
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v3
+      name: Checkout repository
+      with:
+        ref: ${{ env.ref }}
+        fetch-depth: 2
+
+      # We want to enforce the following rules for PRs:
+      # * if all modifications are to README.md
+      #     no testing is needed
+      # * if there are modifications to docs/* or to any code
+      #     then docs need to be built to verify consistency
+      # * if there are modifications to notebooks/* or to any code
+      #     then notebooks need to be run to verify consistency
+      # * for any code changes (or changes to metadata files)
+      #     linting and testing should be run
+      # For a PR build, HEAD will be the merge commit, and we want to diff against the base branch,
+      #  which will be the first parent: HEAD^ 
+      # (For non-PR changes, we will always perform all CI tasks)
+      # Note that GitHub Actions provides path filters, but they operate at the workflow level, not the job level
+    - run: |
+        if ($env:GITHUB_EVENT_NAME -eq 'pull_request') {
+          $editedFiles = git diff HEAD^ --name-only
+          $editedFiles # echo edited files to enable easier debugging
+          $codeChanges = $false
+          $docChanges = $false
+          $nbChanges = $false
+          $changeType = "none"
+          foreach ($file in $editedFiles) {
+            switch -Wildcard ($file) {
+              "README.md" { Continue }
+              ".gitignore" { Continue }
+              "econml/_version.py" { Continue }
+              "prototypes/*" { Continue }
+              "images/*" { Continue }
+              "doc/*" { $docChanges = $true; Continue }
+              "notebooks/*" { $nbChanges = $true; Continue }
+              default { $codeChanges = $true; Continue }
+            }
+          }
+        }
+        echo "buildDocs=$(($env:GITHUB_EVENT_NAME -ne 'pull_request') -or ($docChanges -or $codeChanges))" >> $env:GITHUB_OUTPUT
+        echo "buildNbs=$(($env:GITHUB_EVENT_NAME -ne 'pull_request') -or ($nbChanges -or $codeChanges))" >> $env:GITHUB_OUTPUT
+        echo "testCode=$(($env:GITHUB_EVENT_NAME -ne 'pull_request') -or $codeChanges)" >> $env:GITHUB_OUTPUT
+      shell: pwsh
+      name: Determine type of code change
+      id: eval
+    outputs: 
+      buildDocs: ${{ steps.eval.outputs.buildDocs }}
+      buildNbs: ${{ steps.eval.outputs.buildNbs }}
+      testCode: ${{ steps.eval.outputs.testCode }}
+
+  lint:
+    name: Lint code
+    needs: [eval]
+    if: ${{ needs.eval.outputs.testCode == 'True' }}
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v3
+      name: Checkout repository
+      with:
+        ref: ${{ env.ref }}
+    - uses: actions/setup-python@v4
+      name: Setup Python
+      with:
+        python-version: 3.9
+    - run: python -m pip install --upgrade pip && pip install --upgrade setuptools
+      name: Ensure latest pip and setuptools
+    - run: 'pip install pycodestyle && pycodestyle econml'
+
+  notebooks:
+    name: Run notebooks
+    needs: [eval]
+    if: ${{ needs.eval.outputs.buildNbs == 'True' }}
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        kind: [except customer scenarios, customer scenarios]
+        include:
+        - kind: "except customer scenarios"
+          extras: "[tf,plt]"
+          pattern: "(?!CustomerScenarios)"
+          install_graphviz: true
+          version: 3.8 # no supported version of tensorflow for 3.9
+        - kind: "customer scenarios"
+          extras: "[plt,dowhy]"
+          pattern: "CustomerScenarios"
+          version: 3.9
+          install_graphviz: false
+      fail-fast: false
+    steps:
+    - uses: actions/checkout@v3
+      name: Checkout repository
+      with:
+        ref: ${{ env.ref }}
+    - uses: actions/setup-python@v4
+      name: Setup Python
+      with:
+        python-version: ${{ matrix.version }}
+    - run: python -m pip install --upgrade pip && pip install --upgrade setuptools
+      name: Ensure latest pip and setuptools
+    - run: sudo apt-get -yq install graphviz
+      name: Install graphviz
+      if: ${{ matrix.install_graphviz }}
+    - run: pip install -e .${{ matrix.extras }}
+      name: Install econml
+    - run: pip install pytest pytest-runner jupyter jupyter-client nbconvert nbformat seaborn xgboost tqdm
+      name: Install test and notebook requirements
+    - run: pip list
+      name: List installed packages
+    - run: python setup.py pytest
+      name: Run notebook tests
+      env:
+        PYTEST_ADDOPTS: '-m "notebook"'
+        NOTEBOOK_DIR_PATTERN: ${{ matrix.pattern }}
+    
+  tests:
+    name: "Run tests"
+    needs: [eval]
+    if: ${{ needs.eval.outputs.testCode == 'True' }}
+    strategy:
+      matrix:
+        os: [ubuntu-latest, windows-latest, macos-latest]
+        python-version: [3.6, 3.7, 3.8, 3.9]
+        kind: [serial, other, dml, main, treatment]
+        exclude:
+          # Serial tests fail randomly on mac sometimes, so we don't run them there
+          - os: macos-latest
+            kind: serial
+          # Python 3.6 isn't supported on ubuntu-latest
+          - os: ubuntu-latest
+            python-version: 3.6
+
+        # Assign the correct package and testing options for each kind of test
+        include:
+          - kind: serial
+            opts: '-m "serial" -n 1'
+            extras: "[tf,plt]"
+          - kind: other
+            opts: '-m "cate_api" -n auto'
+            extras: "[tf,plt]"
+          - kind: dml
+            opts: '-m "dml"'
+            extras: "[tf,plt]"
+          - kind: main
+            opts: '-m "not (notebook or automl or dml or serial or cate_api or treatment_featurization)" -n 2'
+            extras: "[tf,plt,dowhy]"
+          - kind: treatment
+            opts: '-m "treatment_featurization" -n auto'
+            extras: "[tf,plt]"
+      fail-fast: false    
+    runs-on: ${{ matrix.os }}
+    steps:
+    - uses: actions/checkout@v3
+      name: Checkout repository
+      with:
+        ref: ${{ env.ref }}
+    - uses: actions/setup-python@v4
+      name: Setup Python
+      with:
+        python-version: ${{ matrix.python-version }}
+    - run: python -m pip install --upgrade pip && pip install --upgrade setuptools
+      name: Ensure latest pip and setuptools
+    - run: pip install -e .${{ matrix.extras }}
+      name: Install econml
+    - run: pip install pytest pytest-runner coverage
+      name: Install pytest
+    - run: python setup.py pytest
+      name: Run tests
+      env:
+        PYTEST_ADDOPTS: ${{ matrix.opts }}
+        COVERAGE_PROCESS_START: 'setup.cfg'
+    # todo: publish test results, coverage info
+      
+  build:
+    name: Build package
+    needs: [eval]
+    if: ${{ needs.eval.outputs.testCode == 'True' }}
+    uses: ./.github/workflows/publish-package.yml
+    with:
+      publish: false
+      repository: testpypi
+      # don't have access to env context here for some reason
+      ref: ${{ github.event_name == 'workflow_dispatch' && inputs.ref || null }}
+
+  docs:
+    name: Build documentation
+    needs: [eval]
+    if: ${{ needs.eval.outputs.buildDocs == 'True' }}
+    uses: ./.github/workflows/publish-documentation.yml
+    with:
+      publish: false
+      environment: test
+      # don't have access to env context here for some reason
+      ref: ${{ github.event_name == 'workflow_dispatch' && inputs.ref || null }}
diff --git a/.github/workflows/publish-documentation.yml b/.github/workflows/publish-documentation.yml
new file mode 100644
index 000000000..dc44f986a
--- /dev/null
+++ b/.github/workflows/publish-documentation.yml
@@ -0,0 +1,80 @@
+name: Build and publish the documentation
+on:
+  workflow_dispatch:
+    inputs:
+      publish:
+        description: 'Whether to publish the documentation (as opposed to just building it)'
+        required: false
+        default: true
+        type: boolean
+      environment:
+        description: 'Whether to publish to production or test environment'
+        required: false
+        default: prod
+        type: choice
+        options: [prod, test]
+      ref:
+        description: 'The git ref to build the documentation for'
+        required: false
+        default: ''
+        type: string
+  # annoyingly, there does not seem to be a way to share these input definitions between triggers
+  workflow_call:
+    inputs:
+      publish:
+        description: 'Whether to publish the documentation (as opposed to just building it)'
+        required: false
+        default: true
+        type: boolean
+      # choice type only supported for workflow_dispatch, not workflow_call
+      environment:
+        description: 'Whether to publish to production or test environment'
+        required: false
+        default: prod
+        type: string
+      ref:
+        description: 'The git ref to build the documentation for'
+        required: false
+        default: ''
+        type: string
+
+
+jobs:
+  create_docs:
+    name: Create and publish documentation
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+        name: Checkout repository
+        with:
+          ref: ${{ inputs.ref }}
+      - uses: actions/setup-python@v4
+        name: Setup Python
+        with:
+          python-version: 3.8 # because of our supported TensorFlow versions, must build on 3.6-3.8
+      - run: python -m pip install --upgrade pip && pip install --upgrade setuptools
+        name: Ensure latest pip and setuptools
+      - run: pip install -U cython && pip install -e .[all]
+        name: Install econml[all]
+      - run: sudo apt-get -yq install graphviz
+        name: Install graphviz
+      - run: pip install "sphinx~=4.4.0" "sphinx_rtd_theme~=1.0.0" && python setup.py build_sphinx -W
+        name: Build documentation
+      - uses: actions/upload-artifact@v3
+        name: Upload docs as artifact
+        with:
+          name: docs
+          path: build/sphinx/html/
+      - run: |-
+          pushd build/sphinx/html
+          zip -r docs.zip *
+          popd
+        name: Zip docs for publishing
+        if : ${{ inputs.publish }}
+      - uses: azure/webapps-deploy@v2
+        name: Deploy documentation to Azure web app
+        with:
+          app-name: ${{ inputs.environment == 'prod' && 'econml' || 'econml-dev' }}
+          package: build/sphinx/html/docs.zip
+          publish-profile: ${{ inputs.environment == 'prod' && secrets.PROD_WEBAPP_PUBLISH_PROFILE || secrets.STAGING_WEBAPP_PUBLISH_PROFILE }}
+        if: ${{ inputs.publish }}          
\ No newline at end of file
diff --git a/.github/workflows/publish-package.yml b/.github/workflows/publish-package.yml
new file mode 100644
index 000000000..6fde706a5
--- /dev/null
+++ b/.github/workflows/publish-package.yml
@@ -0,0 +1,117 @@
+name: Build and publish the package to PyPI
+
+on:
+  workflow_dispatch:
+    inputs:
+      publish:
+        description: 'Whether to publish the package (as opposed to just building it)'
+        required: false
+        default: true
+        type: boolean
+      repository:
+        description: 'Whether to publish to production PyPI or test PyPI'
+        required: false
+        default: pypi
+        type: choice
+        options: [pypi, testpypi]
+      ref:
+        description: 'The git ref to build the package for'
+        required: false
+        default: ''
+        type: string
+  # annoyingly, there does not seem to be a way to share these input definitions between triggers
+  workflow_call:
+    inputs:
+      publish:
+        description: 'Whether to publish the package (as opposed to just building it)'
+        required: false
+        default: true
+        type: boolean
+      # choice type only supported for workflow_dispatch, not workflow_call
+      repository:
+        description: 'Whether to publish to production PyPI or test PyPI'
+        required: false
+        default: pypi
+        type: string
+      ref:
+        description: 'The git ref to build the package for'
+        required: false
+        default: ''
+        type: string
+
+jobs:
+  build_wheels:
+    name: Build wheels on ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-latest, windows-latest, macos-latest]
+    runs-on: ${{ matrix.os }}
+    steps:
+      - uses: actions/checkout@v3
+        name: Checkout repository
+        with:
+          ref: ${{ inputs.ref }}
+      - uses: actions/setup-python@v4
+        name: Setup Python
+        with:
+          python-version: 3.9
+      - run: python -m pip install --upgrade pip && pip install --upgrade setuptools
+        name: Ensure latest pip and setuptools
+      - run: pip install cibuildwheel && python -m cibuildwheel --output-dir dist
+        name: Build wheels
+        env:
+          CIBW_BUILD: cp3{7,8,9,10}-*
+          CIBW_SKIP: "*musl* *win32 *i686"
+      - uses: actions/upload-artifact@v3
+        name: Upload wheels as artifact
+        with:
+          name: dist
+          path: dist/
+  build_sdist:
+    name: Build sdist
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+        name: Checkout repository
+        with:
+          ref: ${{ inputs.ref }}
+      - uses: actions/setup-python@v4
+        name: Setup Python
+        with:
+          python-version: 3.8 # because of our supported TensorFlow versions, must build on 3.6-3.8
+      - run: python -m pip install --upgrade pip && pip install --upgrade setuptools
+        name: Ensure latest pip and setuptools
+      - run: pip install -U cython && pip install -e .[all]
+        name: Install econml[all]
+      - run: python setup.py sdist
+        name: Build sdist
+      - uses: actions/upload-artifact@v3
+        name: Upload sdist as artifact
+        with:
+          name: dist
+          path: dist/
+  publish:
+    name: Publish to PyPI or TestPyPI
+    needs: [build_wheels, build_sdist]
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/setup-python@v4
+        name: Setup Python
+        with:
+          python-version: 3.9
+      - run: python -m pip install --upgrade pip && pip install --upgrade setuptools
+        name: Ensure latest pip and setuptools
+      - run: pip install twine
+        name: Install twine
+      - uses: actions/download-artifact@v3
+        name: Download wheels and sdist
+        with:
+          name: dist
+          path: dist/
+      - run: twine upload dist/*
+        name: Upload wheels and sdist to package index
+        env:
+          TWINE_USERNAME: __token__
+          TWINE_REPOSITORY: ${{ inputs.repository }}
+          TWINE_PASSWORD: ${{ inputs.repository == 'pypi' && secrets.PYPI_UPLOAD_TOKEN  || secrets.TEST_PYPI_UPLOAD_TOKEN }}
+    if: ${{ inputs.publish }}
\ No newline at end of file
diff --git a/azure-pipelines-steps.yml b/azure-pipelines-steps.yml
deleted file mode 100644
index 12f88fa69..000000000
--- a/azure-pipelines-steps.yml
+++ /dev/null
@@ -1,65 +0,0 @@
-# Python package
-# Create and test a Python package on multiple Python versions.
-# Add steps that analyze code, save the dist with the build record, publish to a PyPI-compatible index, and more:
-# https://docs.microsoft.com/azure/devops/pipelines/languages/python
-
-parameters:
-  package: '-e .'
-  images: ['ubuntu-20.04', 'macOS-12', 'windows-2019']
-  versions: ['3.7', '3.8', '3.9', '3.10']
-  job: 
-    job: Job
-  
-jobs:
-- job: ${{ parameters.job.job }}
-  variables:
-    ${{ if ge(length(parameters.images), 2) }}:
-      imgPart: ", {0}"
-    ${{ if lt(length(parameters.images), 2) }}:
-      imgPart: ""
-    ${{ if ge(length(parameters.versions), 2) }}:
-      verPart: ", Python {1}"
-    ${{ if lt(length(parameters.versions), 2) }}:
-      verPart: ""
-
-  strategy:
-    matrix:
-      ${{ each image in parameters.images }}:
-        ${{ each version in parameters.versions }}:
-          ${{ format(format('{0}{1} ', variables.imgPart, variables.verPart), image, version) }}:
-            imageName: ${{ image }}
-            python.version: ${{ version }}
-
-  pool:
-    vmImage: $(imageName)
-
-  ${{ each pair in parameters.job }}: # Insert all properties other than "strategy", "pool", "steps", "variables", "job"
-    ${{ if notIn(pair.key, 'strategy', 'pool', 'steps', 'variables', 'job') }}:
-      ${{ pair.key }}: ${{ pair.value }}
-
-  steps:
-  - task: UsePythonVersion@0
-    displayName: 'Use Python $(python.version)'
-    inputs:
-      versionSpec: '$(python.version)'
-
-  # Enable long path support on Windows so that all packages can be installed correctly
-  - script: 'reg add HKLM\SYSTEM\CurrentControlSet\Control\FileSystem /v LongPathsEnabled /t REG_DWORD /d 1 /f'
-    displayName: 'Enable long paths on Windows'
-    condition: and(succeeded(), eq(variables['Agent.OS'], 'Windows_NT'))
-
-  # Install graphviz programmatically on Linux
-  - script: 'sudo apt-get -yq install graphviz'
-    displayName: 'Install graphviz on Linux'
-    condition: and(succeeded(), eq(variables['Agent.OS'], 'Linux'))
-
-  # Install OpenMP on Mac to support lightgbm
-  - script: 'brew install libomp'
-    displayName: 'Install OpenMP on Mac'
-    condition: and(succeeded(), eq(variables['Agent.OS'], 'Darwin'))
-
-  # Install the package
-  - script: 'python -m pip install --upgrade pip && pip install --upgrade setuptools wheel Cython && pip install ${{ parameters.package }} && pip freeze --exclude-editable'
-    displayName: 'Install dependencies'
-
-  - ${{ parameters.job.steps }}
diff --git a/azure-pipelines.yml b/azure-pipelines.yml
deleted file mode 100644
index 772c405d0..000000000
--- a/azure-pipelines.yml
+++ /dev/null
@@ -1,333 +0,0 @@
-# Python package
-# Create and test a Python package on multiple Python versions.
-# Add steps that analyze code, save the dist with the build record, publish to a PyPI-compatible index, and more:
-# https://docs.microsoft.com/azure/devops/pipelines/languages/python
-
-trigger:
-- main
-
-jobs:
-- job: 'EvalChanges'
-  displayName: 'Analyze changed files to determine which job to run'
-  pool:
-    vmImage: 'macos-12'
-  steps:
-  # We want to enforce the following rules for PRs:
-  # * if all modifications are to README.md
-  #     no testing is needed
-  # * if there are modifications to docs/* or to any code
-  #     then docs need to be built to verify consistency
-  # * if there are modifications to notebooks/* or to any code
-  #     then notebooks need to be run to verify consistency
-  # * for any code changes (or changes to metadata files)
-  #     linting and testing should be run
-  # For a PR build, HEAD will be the merge commit, and we want to diff against the base branch,
-  #  which will be the first parent: HEAD^ 
-  # (For non-PR changes, we will always perform all CI tasks)
-  - powershell: |
-      if ($env:BUILD_REASON -eq 'PullRequest') {
-        $editedFiles = git diff HEAD^ --name-only
-        $editedFiles # echo edited files to enable easier debugging
-        $codeChanges = $false
-        $docChanges = $false
-        $nbChanges = $false
-        $changeType = "none"
-        foreach ($file in $editedFiles) {
-          switch -Wildcard ($file) {
-            "README.md" { Continue }
-            ".gitignore" { Continue }
-            "econml/_version.py" { Continue }
-            "prototypes/*" { Continue }
-            "images/*" { Continue }
-            "doc/*" { $docChanges = $true; Continue }
-            "notebooks/*" { $nbChanges = $true; Continue }
-            default { $codeChanges = $true; Continue }
-          }
-        }
-      }
-      Write-Host "##vso[task.setvariable variable=buildDocs;isOutput=true]$(($env:BUILD_REASON -ne 'PullRequest') -or ($docChanges -or $codeChanges))"
-      Write-Host "##vso[task.setvariable variable=buildNbs;isOutput=true]$(($env:BUILD_REASON -ne 'PullRequest') -or ($nbChanges -or $codeChanges))"
-      Write-Host "##vso[task.setvariable variable=testCode;isOutput=true]$(($env:BUILD_REASON -ne 'PullRequest') -or $codeChanges)"
-    name: output
-    displayName: 'Determine type of code change'
-
-- template: azure-pipelines-steps.yml
-  parameters:
-    versions: ['3.8']
-    images: ['ubuntu-20.04']
-    package: '-e .[all]'
-    job:
-      job: 'Docs'
-      displayName: 'Build documentation'  
-      dependsOn: 'EvalChanges'
-      condition: eq(dependencies.EvalChanges.outputs['output.buildDocs'], 'True')
-      steps:
-      - script: 'sudo apt-get -yq install graphviz'
-        displayName: 'Install graphviz'
-
-      - script: 'pip install sklearn-contrib-lightning'
-        displayName: 'Install lightning'        
-
-      - script: 'pip install git+https://github.com/slundberg/shap.git@d1d2700acc0259f211934373826d5ff71ad514de'
-        displayName: 'Install specific version of shap'  
-
-      - script: 'pip install sphinx!=5.1.0 sphinx_rtd_theme'
-        displayName: 'Install sphinx'
-
-      - script: 'python setup.py build_sphinx -W'
-        displayName: 'Build documentation'
-
-      - publish: 'build/sphinx/html'
-        condition: 'succeededOrFailed()'
-        artifact: 'Documentation'
-        displayName: 'Publish documentation as artifact'
-
-      - script: 'python setup.py build_sphinx -b doctest'
-        displayName: 'Run doctests'
-
-
-- template: azure-pipelines-steps.yml
-  parameters:
-    versions: ['3.8']
-    images: ['ubuntu-20.04']
-    package: '-e .[tf,plt,dowhy]'
-    job:
-      job: 'Notebooks_cust'
-      dependsOn: 'EvalChanges'
-      condition: eq(dependencies.EvalChanges.outputs['output.buildNbs'], 'True')
-      displayName: 'Notebooks (Customer Solutions)'
-      steps:
-      # Work around https://github.com/pypa/pip/issues/9542
-      - script: 'pip install -U numpy~=1.21.0'
-        displayName: 'Upgrade numpy'
-
-      - script: 'pip install pytest pytest-runner jupyter jupyter-client nbconvert nbformat seaborn xgboost tqdm py && pip list && python setup.py pytest'
-        displayName: 'Unit tests'
-        env:
-          PYTEST_ADDOPTS: '-m "notebook"'
-          NOTEBOOK_DIR_PATTERN: 'CustomerScenarios'
-
-      - task: PublishTestResults@2
-        displayName: 'Publish Test Results **/test-results.xml'
-        inputs:
-          testResultsFiles: '**/test-results.xml'
-          testRunTitle: 'Notebooks'
-        condition: succeededOrFailed()
-
-- template: azure-pipelines-steps.yml
-  parameters:
-    versions: ['3.8']
-    images: ['ubuntu-20.04']
-    package: '-e .[tf,plt]'
-    job:
-      job: 'Notebooks_noncust'
-      dependsOn: 'EvalChanges'
-      condition: eq(dependencies.EvalChanges.outputs['output.buildNbs'], 'True')
-      displayName: 'Notebooks (except Customer Solutions)'
-      steps:
-      # Work around https://github.com/pypa/pip/issues/9542
-      - script: 'pip install -U numpy~=1.21.0'
-        displayName: 'Upgrade numpy'
-        
-      - script: 'pip install pytest pytest-runner jupyter jupyter-client nbconvert nbformat seaborn xgboost tqdm py && python setup.py pytest'
-        displayName: 'Unit tests'
-        env:
-          PYTEST_ADDOPTS: '-m "notebook"'
-          NOTEBOOK_DIR_PATTERN: '(?!CustomerScenarios)'
-
-      - task: PublishTestResults@2
-        displayName: 'Publish Test Results **/test-results.xml'
-        inputs:
-          testResultsFiles: '**/test-results.xml'
-          testRunTitle: 'Notebooks'
-        condition: succeededOrFailed()
-
-
-# - job: 'AutoML'
-#   dependsOn: 'EvalChanges'
-#   condition: eq(dependencies.EvalChanges.outputs['output.testCode'], 'True')
-#   variables:
-#     python.version: '3.6'
-#   pool:
-#     vmImage: 'ubuntu-20.04'
-#   steps:
-#   - template: azure-pipelines-steps.yml
-#     parameters:
-#       body:
-#         - task: AzureCLI@2
-#           displayName: 'AutoML tests'
-#           inputs:
-#             azureSubscription: 'automl'
-#             scriptLocation: 'inlineScript'
-#             scriptType: 'pscore'
-#             powerShellIgnoreLASTEXITCODE: '' # string for now due to https://github.com/microsoft/azure-pipelines-tasks/issues/12266
-#             inlineScript: |
-#               $env:SUBSCRIPTION_ID = az account show --query id -o tsv
-#               python setup.py pytest
-#           env:
-#             WORKSPACE_NAME: 'testWorkspace'
-#             RESOURCE_GROUP: 'testingAutoMLEconML'
-#             PYTEST_ADDOPTS: '-m "automl" -n 0'
-#             COVERAGE_PROCESS_START: 'setup.cfg'
-
-#         - task: PublishTestResults@2
-#           displayName: 'Publish Test Results **/test-results.xml'
-#           inputs:
-#             testResultsFiles: '**/test-results.xml'
-#             testRunTitle: 'AutoML'
-#           condition: succeededOrFailed()
-#       package: '.[automl]'
-
-- template: azure-pipelines-steps.yml
-  parameters:
-    versions: ['3.8']
-    images: ['macos-12']
-    job:
-      job: 'Linting'
-      dependsOn: 'EvalChanges'
-      condition: eq(dependencies.EvalChanges.outputs['output.testCode'], 'True')
-      steps:
-      - script: 'pip install pycodestyle && pycodestyle econml'
-        failOnStderr: true
-        displayName: Linting
-
-- template: azure-pipelines-steps.yml
-  parameters:
-    package: '-e .[tf,plt,dowhy]'
-    job:
-      job: Tests_main
-      dependsOn: 'EvalChanges'
-      condition: eq(dependencies.EvalChanges.outputs['output.testCode'], 'True')
-      displayName: 'Run tests (main)'
-      steps:
-      - script: 'pip install pytest pytest-runner "coverage<6.4.1;python_version==''3.6''" "coverage;python_version>''3.6''" py && python setup.py pytest'
-        displayName: 'Unit tests'
-        env:
-          PYTEST_ADDOPTS: '-m "not (notebook or automl or dml or serial or cate_api or treatment_featurization)" -n 2'
-          COVERAGE_PROCESS_START: 'setup.cfg'
-      - task: PublishTestResults@2
-        displayName: 'Publish Test Results **/test-results.xml'
-        inputs:
-          testResultsFiles: '**/test-results.xml'
-          testRunTitle: 'Python $(python.version), image $(imageName)'
-        condition: succeededOrFailed()
-
-      - task: PublishCodeCoverageResults@1
-        displayName: 'Publish Code Coverage Results'
-        inputs:
-          codeCoverageTool: Cobertura
-          summaryFileLocation: '$(System.DefaultWorkingDirectory)/**/coverage.xml'
-
-- template: azure-pipelines-steps.yml
-  parameters:
-    package: '-e .[tf,plt]'
-    job:
-      job: Tests_dml
-      dependsOn: 'EvalChanges'
-      condition: eq(dependencies.EvalChanges.outputs['output.testCode'], 'True')
-      displayName: 'Run tests (DML)'
-      steps:
-      - script: 'pip install pytest pytest-runner "coverage<6.4.1;python_version==''3.6''" "coverage;python_version>''3.6''" py && python setup.py pytest'
-        displayName: 'Unit tests'
-        env:
-          PYTEST_ADDOPTS: '-m "dml"'
-          COVERAGE_PROCESS_START: 'setup.cfg'
-      - task: PublishTestResults@2
-        displayName: 'Publish Test Results **/test-results.xml'
-        inputs:
-          testResultsFiles: '**/test-results.xml'
-          testRunTitle: 'Python $(python.version), image $(imageName)'
-        condition: succeededOrFailed()
-
-      - task: PublishCodeCoverageResults@1
-        displayName: 'Publish Code Coverage Results'
-        inputs:
-          codeCoverageTool: Cobertura
-          summaryFileLocation: '$(System.DefaultWorkingDirectory)/**/coverage.xml'
-
-- template: azure-pipelines-steps.yml
-  parameters:
-    # exclude macOS since these tests frequently break there
-    images: ['ubuntu-20.04', 'windows-2019']
-    package: '-e .[tf,plt]'
-    job:
-      job: Tests_serial
-      dependsOn: 'EvalChanges'
-      condition: eq(dependencies.EvalChanges.outputs['output.testCode'], 'True')
-      displayName: 'Run tests (Serial)'
-      steps:
-      - script: 'pip install pytest pytest-runner "coverage<6.4.1;python_version==''3.6''" "coverage;python_version>''3.6''" py && python setup.py pytest'
-        displayName: 'Unit tests'
-        env:
-          PYTEST_ADDOPTS: '-m "serial" -n 1'
-          COVERAGE_PROCESS_START: 'setup.cfg'
-      - task: PublishTestResults@2
-        displayName: 'Publish Test Results **/test-results.xml'
-        inputs:
-          testResultsFiles: '**/test-results.xml'
-          testRunTitle: 'Python $(python.version), image $(imageName)'
-        condition: succeededOrFailed()
-
-      - task: PublishCodeCoverageResults@1
-        displayName: 'Publish Code Coverage Results'
-        inputs:
-          codeCoverageTool: Cobertura
-          summaryFileLocation: '$(System.DefaultWorkingDirectory)/**/coverage.xml'
-
-- template: azure-pipelines-steps.yml
-  parameters:
-    package: '-e .[tf,plt]'
-    job:
-      job: Tests_CATE_API
-      dependsOn: 'EvalChanges'
-      condition: eq(dependencies.EvalChanges.outputs['output.testCode'], 'True')
-      displayName: 'Run tests (Other)'
-      steps:
-      - script: 'pip install pytest pytest-runner "coverage<6.4.1;python_version==''3.6''" "coverage;python_version>''3.6''" py'
-        displayName: 'Install pytest'
-      - script: 'python setup.py pytest'
-        displayName: 'CATE Unit tests'
-        env:
-          PYTEST_ADDOPTS: '-m "cate_api" -n auto'
-          COVERAGE_PROCESS_START: 'setup.cfg'
-      - task: PublishTestResults@2
-        displayName: 'Publish Test Results **/test-results.xml'
-        inputs:
-          testResultsFiles: '**/test-results.xml'
-          testRunTitle: 'Python $(python.version), image $(imageName)'
-        condition: succeededOrFailed()
-
-      - task: PublishCodeCoverageResults@1
-        displayName: 'Publish Code Coverage Results'
-        inputs:
-          codeCoverageTool: Cobertura
-          summaryFileLocation: '$(System.DefaultWorkingDirectory)/**/coverage.xml'
-
-- template: azure-pipelines-steps.yml
-  parameters:
-    package: '-e .[tf,plt]'
-    job:
-      job: Tests_treatment_featurization
-      dependsOn: 'EvalChanges'
-      condition: eq(dependencies.EvalChanges.outputs['output.testCode'], 'True')
-      displayName: 'Run tests (Treatment Featurization)'
-      steps:
-      - script: 'pip install pytest pytest-runner "coverage<6.4.1;python_version==''3.6''" "coverage;python_version>''3.6''" py'
-        displayName: 'Install pytest'
-      - script: 'python setup.py pytest'
-        displayName: 'Treatment Featurization Unit tests'
-        env:
-          PYTEST_ADDOPTS: '-m "treatment_featurization" -n auto'
-          COVERAGE_PROCESS_START: 'setup.cfg'
-      - task: PublishTestResults@2
-        displayName: 'Publish Test Results **/test-results.xml'
-        inputs:
-          testResultsFiles: '**/test-results.xml'
-          testRunTitle: 'Python $(python.version), image $(imageName)'
-        condition: succeededOrFailed()
-
-      - task: PublishCodeCoverageResults@1
-        displayName: 'Publish Code Coverage Results'
-        inputs:
-          codeCoverageTool: Cobertura
-          summaryFileLocation: '$(System.DefaultWorkingDirectory)/**/coverage.xml'
\ No newline at end of file
diff --git a/setup.cfg b/setup.cfg
index eb1d53122..62107f18e 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -44,8 +44,8 @@ install_requires =
 test_suite = econml.tests
 tests_require =
     pytest
-    pytest-xdist < 2.0.0
-    pytest-cov < 3.0.0
+    pytest-xdist
+    pytest-cov
 
 [options.extras_require]
 automl =