StatusIs(error::INTERNAL,
                       HasSubstr(absl::StrCat("PjRtClient for ", DEVICE_CPU,
                                              " is not type PjRtCApiClient"))));
}

TEST(TensorPjRtBufferUtilTest, GetPjRtCApiClientSuccess) {
  auto status = pjrt::PjrtApi(DEVICE_CPU);
  if (!status.ok()) {
    TF_ASSERT_OK(pjrt::SetPjrtApi(DEVICE_CPU, GetPjrtApi()));
  }

  DeviceType device_gpu = DeviceType(DEVICE_GPU);
  EXPECT_EQ(GetPersistentCacheDirectory(device_gpu), "/tmp/xla_cache");
  DeviceType device_cpu = DeviceType(DEVICE_CPU);
  EXPECT_EQ(GetPersistentCacheDirectory(device_cpu), "/tmp/xla_cache");
  DeviceType device_tpu = DeviceType(DEVICE_TPU);
  EXPECT_TRUE(GetPersistentCacheDirectory(device_tpu).empty());
}

  int i = 0;
  if (ctx->has_input(i) || ctx->has_input(++i)) {
    ctx->set_output(0, ctx->input(i));
  }
}

REGISTER_KERNEL_BUILDER(Name("XlaLaunch").Device(DEVICE_CPU), XlaLocalLaunchOp);

REGISTER_KERNEL_BUILDER(Name("XlaLaunchV2").Device(DEVICE_CPU), XlaLaunchV2Op);

REGISTER_KERNEL_BUILDER(Name("XlaLaunch")
                            .Device(DEVICE_GPU)
                            .HostMemory("constants")

          input_flat.data()[i], lower_bound, bin_width, num_bins);
      histogram_flat.data()[bin_index] += 1;
    }
  }
};

REGISTER_KERNEL_BUILDER(Name("CustomAggregator").Device(DEVICE_CPU),
                        CustomAggregatorOp);

  }

 private:
  bool enabled_;
  std::unique_ptr<tsl::WritableFile> tensor_data_file_;
  std::unique_ptr<io::RecordWriter> tensor_data_writer_;
};

REGISTER_KERNEL_BUILDER(Name("DumpTensor").Device(DEVICE_CPU), DumpTensorOp);

  // CPU devices of the same job as the TPU_SYSTEM device.
  system_device.type = tensorflow::DEVICE_CPU;
  return DeviceNameUtils::ParsedNameToString(system_device);
}

// Find the host CPU device for a given TPU device with `DEVICE_CPU` as its
// type. If multiple local cpu devices are disabled, always assign id 0. If
// set, use the same id as the tpu device.

           stage == IrExportStage::HLO_SERIALIZED;
  };
  // TODO(b/238830423): support GetCompilerIr on TFRT TPU device for stages
  // that requires compilation from HLO to executable.
  if (device_type != DEVICE_CPU && stream == nullptr &&
      !is_tfrt_tpu_supported_stage(stage)) {
    return absl::InternalError(
        "GetCompilerIr with requested stage is not supported on this device.");
  }
  return absl::OkStatus();
}

      default:
        collector = std::make_unique<CalibrationStatisticsCollectorMinMax>();
    }
  }
};

REGISTER_KERNEL_BUILDER(Name("CalibrationStatisticsSaver").Device(DEVICE_CPU),
                        CalibrationStatisticsSaverOp);

  DeviceType* device_type = id_to_device_type_.back().get();
  TF_RETURN_IF_ERROR(DeviceNameToDeviceType(names_.back(), device_type));

  is_cpu_.push_back(device_type->type_string() == DEVICE_CPU);
  is_gpu_.push_back(device_type->type_string() == DEVICE_GPU);

  name_to_id_.emplace(string(name), DeviceId(new_id));

  const XlaOpRegistry::DeviceRegistration* compilation_device;

    // Early return for any node with a device that is not a CPU or GPU.
    DeviceNameUtils::ParsedName parsed;
    if (DeviceNameUtils::ParseFullName(n->requested_device(), &parsed)) {
      if (parsed.type != DEVICE_CPU && parsed.type != DEVICE_GPU) {
        return false;
      }
    }
  }
  return true;
}

// Checks if a graph node is marked to be a guaranteed constant.
bool is_guaranteed_constant(const Node& n) {