Class PyLocalExecutable

tensorflow/compiler/xla/python/local_client.h:250–354

Represents a compiled computation that can be executed given handles to device-allocated literals. Wraps one or more XLA LocalExecutables (one per partition, as specified by the build options).

// device-allocated literals. Wraps one or more XLA LocalExecutables (one per
// partition, as specified by the build options).
class PyLocalExecutable {
 public:
  // Compiles a computation to an executable.
  static StatusOr<std::unique_ptr<PyLocalExecutable>> CompileForDevices(
      const XlaComputation& computation,
      absl::optional<std::vector<Shape>> argument_layouts,
      const ExecutableBuildOptions* build_options,
      std::shared_ptr<PyLocalClient> client,
      const std::vector<std::vector<std::shared_ptr<Device>>>&
          device_assignment);

  // TODO(phawkins): Deprecated. Delete once all callers have been updated to
  // use the newer form.
  static StatusOr<std::unique_ptr<PyLocalExecutable>> Compile(
      const XlaComputation& computation,
      absl::optional<std::vector<Shape>> argument_layouts,
      const ExecutableBuildOptions* build_options,
      std::shared_ptr<PyLocalClient> client,
      absl::optional<DeviceAssignment> device_assignment);

  PyLocalExecutable(std::vector<std::unique_ptr<LocalExecutable>> executables,
                    DeviceAssignment device_assignment,
                    std::shared_ptr<PyLocalClient> client);

  int num_replicas() const {
    return executables_[0]->build_options().num_replicas();
  }

  int num_partitions() const {
    return executables_[0]->build_options().num_partitions();
  }

  int64 SizeOfGeneratedCodeInBytes() const {
    int64 size = 0;
    for (auto& executable : executables_) {
      size += executable->executable()->SizeOfGeneratedCodeInBytes();
    }
    return size;
  }

  const std::vector<std::shared_ptr<LocalExecutable>>& executables() const {
    return executables_;
  }

  const DeviceAssignment& device_assignment() const {
    return *device_assignment_;
  }

  const std::vector<std::pair<int, int>>& local_logical_device_ids() const {
    return local_logical_device_ids_;
  }

  const std::vector<std::shared_ptr<Device>>& local_devices() const {
    return local_devices_;
  }

  StatusOr<std::unique_ptr<PyLocalBuffer>> Execute(
      absl::Span<PyLocalBuffer* const> argument_handles);