//
// Use in tests to allow XLA to fallback to TF classic. This has global effect.
TF_CAPI_EXPORT unsigned char TF_SetXlaEnableLazyCompilation(
    unsigned char enable);
TF_CAPI_EXPORT unsigned char TF_SetTfXlaCpuGlobalJit(unsigned char enable);

// Sets XLA's auto jit mode according to the specified string, which is parsed
// as if passed in XLA_FLAGS. This has global effect.
TF_CAPI_EXPORT void TF_SetXlaAutoJitMode(const char* mode);

// Sets a thread-local device placement policy. After this call, other calls to
// TFE_Execute in the same thread will use the device policy specified here
// instead of the device policy used to construct the context. This has no
// effect on the device policy used by other program threads.
TF_CAPI_EXPORT extern void TFE_ContextSetThreadLocalDevicePlacementPolicy(
    TFE_Context* ctx, TFE_ContextDevicePlacementPolicy policy);

  // Fill the cache.
  TF_EXPECT_OK(ReadCache(&cache, "a", 0, n, &out));
  EXPECT_EQ(out, a);
  EXPECT_EQ(calls, 1);
  TF_EXPECT_OK(ReadCache(&cache, "a", 8, n, &out));
  EXPECT_EQ(out, A);
  EXPECT_EQ(calls, 2);
  TF_EXPECT_OK(ReadCache(&cache, "b", 0, n, &out));
  EXPECT_EQ(out, b);
  EXPECT_EQ(calls, 3);
  TF_EXPECT_OK(ReadCache(&cache, "b", 8, n, &out));
  EXPECT_EQ(out, B);
  EXPECT_EQ(calls, 4);

  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
  float product[4] = {0};
  EXPECT_EQ(sizeof(product), TF_TensorByteSize(t));
  memcpy(&product[0], TF_TensorData(t), TF_TensorByteSize(t));
  TF_DeleteTensor(t);
  EXPECT_EQ(7, product[0]);
  EXPECT_EQ(10, product[1]);
  EXPECT_EQ(15, product[2]);
  EXPECT_EQ(22, product[3]);
  TF_DeleteStatus(status);
}

  ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());
  TF_Tensor* result_tensor = TFE_TensorHandleResolve(result_t, status.get());
  float* result_value = static_cast<float*>(TF_TensorData(result_tensor));
  EXPECT_EQ(*result_value, 4.0);

  TF_DeleteTensor(result_tensor);
  TF_DeleteAbstractTensor(result);
  TF_DeleteOutputList(o);
  TF_DeleteExecutionContext(ctx);
}

// MatMul Test

// in_flight_nodes_limit: when is_async is true, this value controls the
// maximum number of in flight async nodes. Enqueuing of additional async ops
// after the limit is reached blocks until some inflight nodes finishes.
// The effect is bounding the memory held by inflight TensorHandles that are
// referenced by the inflight nodes.
// A recommended value has not been established.
// A value of 0 removes the limit, which is the behavior of TensorFlow 2.11.

  ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
  EXPECT_EQ(TFE_TensorHandleDataType(packed_handle), TF_RESOURCE);
  EXPECT_EQ(TFE_TensorHandleNumDims(packed_handle, status), 0);
  EXPECT_EQ(TFE_TensorHandleNumElements(packed_handle, status), 1);

  const string composite_device_name =
      "/job:localhost/replica:0/task:0/device:COMPOSITE:0";
  EXPECT_EQ(TFE_TensorHandleDeviceName(packed_handle, status),

bool GetGraphDef(TF_Graph* graph, tensorflow::GraphDef* graph_def) {
  TF_Status* s = TF_NewStatus();
  TF_Buffer* buffer = TF_NewBuffer();
  TF_GraphToGraphDef(graph, buffer, s);
  bool ret = TF_GetCode(s) == TF_OK;
  EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);
  if (ret) ret = graph_def->ParseFromArray(buffer->data, buffer->length);
  TF_DeleteBuffer(buffer);
  TF_DeleteStatus(s);
  return ret;
}

  EXPECT_FALSE(executed);
  ASSERT_FALSE(TF_GetCode(status.get()) == TF_OK);

  // Custom devices can refuse to do type-based dispatch (as hcustom1 is
  // configured to do)
  matmul.reset(MatMulOp(context.get(), hcustom1.get(), hcpu.get()));
  num_retvals = 1;
  executed = false;
  TFE_Execute(matmul.get(), &retval, &num_retvals, status.get());
  EXPECT_FALSE(executed);

    int64_t* dims = output_shapes->items[0].dims;

    if (!expected_shape.has_value()) {
      EXPECT_EQ(num_dims, -1);
      EXPECT_EQ(dims, nullptr);
      return;
    }

    EXPECT_EQ(num_dims, expected_shape->size());
    for (size_t i = 0; i < num_dims; ++i) {
      EXPECT_EQ(dims[i], (*expected_shape)[i]);
    }
    TF_DeleteShapeAndTypeList(input_shapes);