service

type BufferAllocationProto struct {
	Index                       int64                             `protobuf:"varint,1,opt,name=index" json:"index,omitempty"`
	Size                        int64                             `protobuf:"varint,2,opt,name=size" json:"size,omitempty"`
	IsThreadLocal               bool                              `protobuf:"varint,3,opt,name=is_thread_local,json=isThreadLocal" json:"is_thread_local,omitempty"`
	IsReusable                  bool                              `protobuf:"varint,4,opt,name=is_reusable,json=isReusable" json:"is_reusable,omitempty"`
	IsEntryComputationParameter bool                              `` /* 140-byte string literal not displayed */
	ParameterNumber             int64                             `protobuf:"varint,6,opt,name=parameter_number,json=parameterNumber" json:"parameter_number,omitempty"`
	MaybeLiveOut                bool                              `protobuf:"varint,7,opt,name=maybe_live_out,json=maybeLiveOut" json:"maybe_live_out,omitempty"`
	Color                       int64                             `protobuf:"varint,8,opt,name=color" json:"color,omitempty"`
	Assigned                    []*BufferAllocationProto_Assigned `protobuf:"bytes,9,rep,name=assigned" json:"assigned,omitempty"`
}

type BufferAllocationProto_Assigned struct {
	LogicalBufferId int64 `protobuf:"varint,1,opt,name=logical_buffer_id,json=logicalBufferId" json:"logical_buffer_id,omitempty"`
	Offset          int64 `protobuf:"varint,2,opt,name=offset" json:"offset,omitempty"`
	Size            int64 `protobuf:"varint,3,opt,name=size" json:"size,omitempty"`
}

type BufferAssignmentProto struct {
	LogicalBuffers      []*LogicalBufferProto                `protobuf:"bytes,1,rep,name=logical_buffers,json=logicalBuffers" json:"logical_buffers,omitempty"`
	BufferAliases       []*BufferAssignmentProto_BufferAlias `protobuf:"bytes,2,rep,name=buffer_aliases,json=bufferAliases" json:"buffer_aliases,omitempty"`
	BufferAllocations   []*BufferAllocationProto             `protobuf:"bytes,3,rep,name=buffer_allocations,json=bufferAllocations" json:"buffer_allocations,omitempty"`
	HeapSimulatorTraces []*HeapSimulatorTrace                `protobuf:"bytes,4,rep,name=heap_simulator_traces,json=heapSimulatorTraces" json:"heap_simulator_traces,omitempty"`
}

type BufferAssignmentProto_BufferAlias struct {
	SourceBufferId int64                        `protobuf:"varint,1,opt,name=source_buffer_id,json=sourceBufferId" json:"source_buffer_id,omitempty"`
	Location       *LogicalBufferProto_Location `protobuf:"bytes,2,opt,name=location" json:"location,omitempty"`
}

type HeapSimulatorTrace struct {
	Events                []*HeapSimulatorTrace_Event `protobuf:"bytes,1,rep,name=events" json:"events,omitempty"`
	WholeModuleSimulation bool                        `protobuf:"varint,2,opt,name=whole_module_simulation,json=wholeModuleSimulation" json:"whole_module_simulation,omitempty"`
}

type HeapSimulatorTrace_Event struct {
	Kind HeapSimulatorTrace_Event_Kind `protobuf:"varint,1,opt,name=kind,enum=xla.HeapSimulatorTrace_Event_Kind" json:"kind,omitempty"`
	// The id of the LogicalBuffer that the event applies to.
	BufferId int64 `protobuf:"varint,2,opt,name=buffer_id,json=bufferId" json:"buffer_id,omitempty"`
	// The HloInstruction that the simulation was processing that caused this
	// event to occur, identified by its computation and instruction name. E.g.
	// buffers defined by instruction A are allocated when processing A.
	ComputationName string `protobuf:"bytes,3,opt,name=computation_name,json=computationName" json:"computation_name,omitempty"`
	InstructionName string `protobuf:"bytes,4,opt,name=instruction_name,json=instructionName" json:"instruction_name,omitempty"`
	// The id of the canonical LogicalBuffer that the buffer shares with. Only
	// set for SHARE_WITH events.
	ShareWithCanonicalId int64 `protobuf:"varint,5,opt,name=share_with_canonical_id,json=shareWithCanonicalId" json:"share_with_canonical_id,omitempty"`
}

type HeapSimulatorTrace_Event_Kind int32

type HloComputationProto struct {
	Name string `protobuf:"bytes,1,opt,name=name" json:"name,omitempty"`
	// The array of instructions is always in a valid dependency order, where
	// operands appear before their users.
	Instructions []*HloInstructionProto `protobuf:"bytes,2,rep,name=instructions" json:"instructions,omitempty"`
}

type HloInstructionProto struct {
	Name                    string          `protobuf:"bytes,1,opt,name=name" json:"name,omitempty"`
	Opcode                  string          `protobuf:"bytes,2,opt,name=opcode" json:"opcode,omitempty"`
	Shape                   *xla.Shape      `protobuf:"bytes,3,opt,name=shape" json:"shape,omitempty"`
	OperandNames            []string        `protobuf:"bytes,4,rep,name=operand_names,json=operandNames" json:"operand_names,omitempty"`
	ControlPredecessorNames []string        `protobuf:"bytes,5,rep,name=control_predecessor_names,json=controlPredecessorNames" json:"control_predecessor_names,omitempty"`
	CalledComputationNames  []string        `protobuf:"bytes,6,rep,name=called_computation_names,json=calledComputationNames" json:"called_computation_names,omitempty"`
	Metadata                *xla.OpMetadata `protobuf:"bytes,7,opt,name=metadata" json:"metadata,omitempty"`
	// Literal, only present for kConstant.
	Literal *xla.LiteralProto `protobuf:"bytes,8,opt,name=literal" json:"literal,omitempty"`
	// Parameter info, only present for kParameter.
	ParameterNumber int64  `protobuf:"varint,9,opt,name=parameter_number,json=parameterNumber" json:"parameter_number,omitempty"`
	ParameterName   string `protobuf:"bytes,10,opt,name=parameter_name,json=parameterName" json:"parameter_name,omitempty"`
	// Fusion state, only present for kFusion.
	FusionKind                   string               `protobuf:"bytes,11,opt,name=fusion_kind,json=fusionKind" json:"fusion_kind,omitempty"`
	FusedInstructionsComputation *HloComputationProto `` /* 141-byte string literal not displayed */
	// Index for kGetTupleElement.
	TupleIndex int64 `protobuf:"varint,13,opt,name=tuple_index,json=tupleIndex" json:"tuple_index,omitempty"`
}

type HloModuleProto struct {
	Name                 string `protobuf:"bytes,1,opt,name=name" json:"name,omitempty"`
	EntryComputationName string `protobuf:"bytes,2,opt,name=entry_computation_name,json=entryComputationName" json:"entry_computation_name,omitempty"`
	// The array of computations is always in a valid dependency order, where
	// callees appear before their callers.
	Computations []*HloComputationProto `protobuf:"bytes,3,rep,name=computations" json:"computations,omitempty"`
}

type HloOrderingProto struct {
	SequentialComputations []*HloOrderingProto_SequentialComputation `protobuf:"bytes,1,rep,name=sequential_computations,json=sequentialComputations" json:"sequential_computations,omitempty"`
}

type HloOrderingProto_SequentialComputation struct {
	ComputationName  string   `protobuf:"bytes,1,opt,name=computation_name,json=computationName" json:"computation_name,omitempty"`
	InstructionNames []string `protobuf:"bytes,2,rep,name=instruction_names,json=instructionNames" json:"instruction_names,omitempty"`
}

type HloProto struct {
	HloModule        *HloModuleProto        `protobuf:"bytes,1,opt,name=hlo_module,json=hloModule" json:"hlo_module,omitempty"`
	HloOrdering      *HloOrderingProto      `protobuf:"bytes,2,opt,name=hlo_ordering,json=hloOrdering" json:"hlo_ordering,omitempty"`
	BufferAssignment *BufferAssignmentProto `protobuf:"bytes,3,opt,name=buffer_assignment,json=bufferAssignment" json:"buffer_assignment,omitempty"`
}

type LogicalBufferProto struct {
	Id   int64 `protobuf:"varint,1,opt,name=id" json:"id,omitempty"`
	Size int64 `protobuf:"varint,2,opt,name=size" json:"size,omitempty"`
	// The location where the buffer is defined.
	DefinedAt *LogicalBufferProto_Location `protobuf:"bytes,3,opt,name=defined_at,json=definedAt" json:"defined_at,omitempty"`
	Color     int64                        `protobuf:"varint,4,opt,name=color" json:"color,omitempty"`
}

type LogicalBufferProto_Location struct {
	// NOTE: module_name isn't necessary, since all LogicalBuffers are
	// associated with a single HloModule.
	ComputationName string  `protobuf:"bytes,1,opt,name=computation_name,json=computationName" json:"computation_name,omitempty"`
	InstructionName string  `protobuf:"bytes,2,opt,name=instruction_name,json=instructionName" json:"instruction_name,omitempty"`
	ShapeIndex      []int64 `protobuf:"varint,3,rep,packed,name=shape_index,json=shapeIndex" json:"shape_index,omitempty"`
}

type OperationRequest struct {
	OutputHandle *xla.ComputationDataHandle `protobuf:"bytes,1,opt,name=output_handle,json=outputHandle" json:"output_handle,omitempty"`
	OutputShape  *xla.Shape                 `protobuf:"bytes,2,opt,name=output_shape,json=outputShape" json:"output_shape,omitempty"`
	// For operations which call embedded computations such as "Map", these are
	// the version(s) that the embedded computation should be called at. A version
	// value of a computation is the ComputationDataHandle of the root of the
	// computation at the point in time.
	//
	// "Call", "Map", "Reduce", and "ReduceWindow" operations take a single
	// embedded computation so this field will have a single value for those
	// operations.
	//
	// "While" operation takes two; index 0 is the "condition" version and index 1
	// is the "body" version.
	EmbeddedComputationVersions []int64 `` /* 145-byte string literal not displayed */
	// The actual request, which in itself is a tagged union of all possible
	// operation request types.
	Request *xla.OpRequest `protobuf:"bytes,4,opt,name=request" json:"request,omitempty"`
}

type SessionComputation struct {
	Name string `protobuf:"bytes,1,opt,name=name" json:"name,omitempty"`
	// The ComputationHandle used to refer to this computation in the XLA
	// service.
	ComputationHandle *xla.ComputationHandle `protobuf:"bytes,2,opt,name=computation_handle,json=computationHandle" json:"computation_handle,omitempty"`
	// Map from ComputationDataHandle value to operation request. The highest
	// ComputationDataHandle value corresponds to the root of the computation.
	Requests map[int64]*OperationRequest `` /* 137-byte string literal not displayed */
}

type SessionModule struct {
	// The entry computation, which was requested for serialization. This may have
	// referred to embedded computations, which are reflected below.
	Entry *SessionComputation `protobuf:"bytes,1,opt,name=entry" json:"entry,omitempty"`
	// Embedded computations that are transitively referred to by the entry
	// computation.
	EmbeddedComputations []*SessionComputation `protobuf:"bytes,2,rep,name=embedded_computations,json=embeddedComputations" json:"embedded_computations,omitempty"`
	// The arguments passed to the computation.
	Arguments []*xla.LiteralProto `protobuf:"bytes,3,rep,name=arguments" json:"arguments,omitempty"`
	// The result of the computation.
	Result *xla.LiteralProto `protobuf:"bytes,4,opt,name=result" json:"result,omitempty"`
	// The name of the platform used to run the computation.
	ExecutionPlatform string `protobuf:"bytes,5,opt,name=execution_platform,json=executionPlatform" json:"execution_platform,omitempty"`
}