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
  # command to reattempt build a few times in the case of failure (b/302558736)
  set +e
  for i in $(seq 1 5)
  do
    docker build \
    --build-arg REQUIREMENTS_FILE=jax.requirements.txt \
    --target=$target \
    --cache-from "$IMAGE" \

  TFE_OpSetAttrType(op.get(), "dtype", type);
  TFE_OpSetAttrShape(op.get(), "shape", dims, num_dims, status);
  TFE_OpSetAttrString(op.get(), "container", "", 0);
  // Use the special GUID for no buffer sharing
  //
  // TODO(allenl): Should we provide a better API for this? AFAIK this is the
  // only reasonable way to make variables with no aliasing using the eager C
  // API.
  std::string no_sharing = "cd2c89b7-88b7-44c8-ad83-06c2a9158347";

  std::vector<std::string> summarized_devices = device_.SummarizeDeviceNames();
  for (int component_index = 0; component_index < tensors_.size();
       ++component_index) {
    // TODO(allenl): Add a C API for summarizing tensors. Currently custom
    // devices limiting themselves to a C API (for ABI compatibility) would need
    // to implement summarization for component tensors themselves.
    ImmediateExecutionTensorHandle* component =

TFE_Op* IdentityOp(TFE_Context* ctx, TFE_TensorHandle* a);

// Return a shape op fetching the shape of `a`.
TFE_Op* ShapeOp(TFE_Context* ctx, TFE_TensorHandle* a);

// Return an allreduce op adding up input tensor `in` from `group_size` workers.
TFE_Op* AllReduceOp(TFE_Context* ctx, TFE_TensorHandle* in, int group_size);

// Return a SendOp op `op_name` with send input tensor `in` and attributes

    return nullptr;
  else
    return &tmp_dir;
}

namespace tensorflow {
namespace {

// TODO(vnvo2409): Refactor `gcs_filesystem_test` to remove unnecessary tests
// after porting all tests from
// `//tensorflow/core/platform/cloud:gcs_file_system_test`.
class GCSFilesystemTest : public ::testing::Test {
 public:
  void SetUp() override {
    root_dir_ = io::JoinPath(
        *GetTmpDir(),

    ASSERT_TRUE(TF_GetCode(status.get()) == TF_INVALID_ARGUMENT)
        << TF_Message(status.get());
  }

  {
    // Try to extract the wrong number of components from a parallel tensor
    std::array<TFE_TensorHandle*, 1> correct_components{value_one.get()};
    TensorHandlePtr combined_value = CreatePerDeviceValues(
        context.get(), correct_components, device_name, status.get());