rm -rf /tf/venv
}

@test "Wheel is manylinux2014 (manylinux_2_17) compliant" {
    python3 -m auditwheel show "$TF_WHEEL" > audit.txt
    grep --quiet -zoP 'is consistent with the following platform tag:\n"manylinux_2_17_(aarch|x86_)64"\.' audit.txt
}

@test "Wheel conforms to upstream size limitations" {
    WHEEL_MEGABYTES=$(stat --format %s "$TF_WHEEL" | awk '{print int($1/(1024*1024))}')
    # Googlers: search for "test_tf_whl_size"

  Args:
    environ_cp: copy of the os.environ.
  """
  if is_ppc64le():
    # gcc on ppc64le does not support -march, use mcpu instead
    default_cc_opt_flags = '-mcpu=native'
  elif is_windows():
    default_cc_opt_flags = '/arch:AVX'
  else:
    # On all other platforms, no longer use `-march=native` as this can result
    # in instructions that are too modern being generated. Users that want

# Any request that includes "nightly_upload" should just use the
# local multi-cache (public read-only cache + disk cache) instead.
export TFCI="$(echo $TFCI | sed 's/,nightly_upload/,public_cache,disk_cache/')"
if [[ -n "${TF_ANY_EXTRA_ENV:-}" ]]; then
  export TFCI="$TFCI,$TF_ANY_EXTRA_ENV"
fi
if [[ -n "${TF_ANY_SCRIPT:-}" ]]; then
  "$TF_ANY_SCRIPT"
elif [[ -n "${TF_ANY_TARGETS:-}" ]]; then

TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1

# 3. Push results to public build cache (unique cache for MacOS)
if [[ $(uname -s) == "Darwin" ]]; then
  TFCI_BAZEL_COMMON_ARGS="$TFCI_BAZEL_COMMON_ARGS --config tf_public_macos_cache_push"
else
  TFCI_BAZEL_COMMON_ARGS="$TFCI_BAZEL_COMMON_ARGS --config tf_public_cache_push"

            gcr.io/tensorflow-sigs/build:${{ steps.vars.outputs.REF }}-${{ matrix.python-version }}
          cache-from: type=registry,ref=tensorflow/build:${{ steps.vars.outputs.REF }}-${{ matrix.python-version }}
          cache-to: type=inline
      -
        name: Image digest

  else
    TAG="pr-${KOKORO_GITHUB_PULL_REQUEST_NUMBER}"
  fi
fi

# Build for both JAX and TF usage.  We do these in one place because they share
# almost all of the same cache layers
export DOCKER_BUILDKIT=1
for target in jax tf; do
  IMAGE="gcr.io/tensorflow-sigs/build-arm64:$target-$TAG"
  docker pull "$IMAGE" || true
  # Due to some flakiness of resources pulled in the build, allow the docker