pjrt_ifrt

package
v2.14.1 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Dec 8, 2023 License: BSD-3-Clause Imports: 8 Imported by: 0

Documentation

Index

Constants

This section is empty.

Variables

View Source 
var File_tensorflow_compiler_xla_python_pjrt_ifrt_xla_compiler_proto protoreflect.FileDescriptor
View Source 
var File_tensorflow_compiler_xla_python_pjrt_ifrt_xla_host_callback_proto protoreflect.FileDescriptor
View Source 
var File_tensorflow_compiler_xla_python_pjrt_ifrt_xla_sharding_proto protoreflect.FileDescriptor

Functions

This section is empty.

Types

type HloShardingProto 

type HloShardingProto struct {
	Devices       *ifrt.DeviceListProto `protobuf:"bytes,1,opt,name=devices,proto3" json:"devices,omitempty"`
	MemoryKind    *string               `protobuf:"bytes,3,opt,name=memory_kind,json=memoryKind,proto3,oneof" json:"memory_kind,omitempty"`
	XlaOpSharding *data.OpSharding      `protobuf:"bytes,2,opt,name=xla_op_sharding,json=xlaOpSharding,proto3" json:"xla_op_sharding,omitempty"`
	// contains filtered or unexported fields
}

Wire format for `HloSharding`.

func (*HloShardingProto) Descriptor deprecated 

func (*HloShardingProto) Descriptor() ([]byte, []int)

Deprecated: Use HloShardingProto.ProtoReflect.Descriptor instead.

func (*HloShardingProto) GetDevices 

func (x *HloShardingProto) GetDevices() *ifrt.DeviceListProto

func (*HloShardingProto) GetMemoryKind 

func (x *HloShardingProto) GetMemoryKind() string

func (*HloShardingProto) GetXlaOpSharding 

func (x *HloShardingProto) GetXlaOpSharding() *data.OpSharding

func (*HloShardingProto) ProtoMessage 

func (*HloShardingProto) ProtoMessage()

func (*HloShardingProto) ProtoReflect 

func (x *HloShardingProto) ProtoReflect() protoreflect.Message

func (*HloShardingProto) Reset 

func (x *HloShardingProto) Reset()

func (*HloShardingProto) String 

func (x *HloShardingProto) String() string

type XlaCompileOptionsProto 

type XlaCompileOptionsProto struct {
	CompileOptions *pjrt.CompileOptionsProto `protobuf:"bytes,1,opt,name=compile_options,json=compileOptions,proto3" json:"compile_options,omitempty"`
	// contains filtered or unexported fields
}

func (*XlaCompileOptionsProto) Descriptor deprecated 

func (*XlaCompileOptionsProto) Descriptor() ([]byte, []int)

Deprecated: Use XlaCompileOptionsProto.ProtoReflect.Descriptor instead.

func (*XlaCompileOptionsProto) GetCompileOptions 

func (x *XlaCompileOptionsProto) GetCompileOptions() *pjrt.CompileOptionsProto

func (*XlaCompileOptionsProto) ProtoMessage 

func (*XlaCompileOptionsProto) ProtoMessage()

func (*XlaCompileOptionsProto) ProtoReflect 

func (x *XlaCompileOptionsProto) ProtoReflect() protoreflect.Message

func (*XlaCompileOptionsProto) Reset 

func (x *XlaCompileOptionsProto) Reset()

func (*XlaCompileOptionsProto) String 

func (x *XlaCompileOptionsProto) String() string

type XlaHostCallbackProto 

type XlaHostCallbackProto struct {

	// The metadata (e.g. channel_id, shape) for the operands and results.
	Operands []*XlaHostCallbackProto_ArgInfo `protobuf:"bytes,1,rep,name=operands,proto3" json:"operands,omitempty"`
	Results  []*XlaHostCallbackProto_ArgInfo `protobuf:"bytes,2,rep,name=results,proto3" json:"results,omitempty"`
	// Serialized host callback.
	SerializedCallback *anypb.Any `protobuf:"bytes,3,opt,name=serialized_callback,json=serializedCallback,proto3" json:"serialized_callback,omitempty"`
	// See comment for PJRT
	// ExecuteOptions::use_major_to_minor_data_layout_for_callbacks.
	UseMajorToMinorDataLayoutForCallbacks bool `` /* 185-byte string literal not displayed */
	// contains filtered or unexported fields
}

Represents a host callback in an XLA computation.

XLA computation may use XLA send/recv to represent communication between the host and the device. This can be used to implement "host callbacks", where host-side computation is invoked in the middle of XLA computation. This message contains information that is necessary to instantiate host callbacks that are marshalled from the client.

Modeled after `xla::HostCallback`.

func (*XlaHostCallbackProto) Descriptor deprecated 

func (*XlaHostCallbackProto) Descriptor() ([]byte, []int)

Deprecated: Use XlaHostCallbackProto.ProtoReflect.Descriptor instead.

func (*XlaHostCallbackProto) GetOperands 

func (x *XlaHostCallbackProto) GetOperands() []*XlaHostCallbackProto_ArgInfo

func (*XlaHostCallbackProto) GetResults 

func (x *XlaHostCallbackProto) GetResults() []*XlaHostCallbackProto_ArgInfo

func (*XlaHostCallbackProto) GetSerializedCallback 

func (x *XlaHostCallbackProto) GetSerializedCallback() *anypb.Any

func (*XlaHostCallbackProto) GetUseMajorToMinorDataLayoutForCallbacks 

func (x *XlaHostCallbackProto) GetUseMajorToMinorDataLayoutForCallbacks() bool

func (*XlaHostCallbackProto) ProtoMessage 

func (*XlaHostCallbackProto) ProtoMessage()

func (*XlaHostCallbackProto) ProtoReflect 

func (x *XlaHostCallbackProto) ProtoReflect() protoreflect.Message

func (*XlaHostCallbackProto) Reset 

func (x *XlaHostCallbackProto) Reset()

func (*XlaHostCallbackProto) String 

func (x *XlaHostCallbackProto) String() string

type XlaHostCallbackProto_ArgInfo 

type XlaHostCallbackProto_ArgInfo struct {

	// The channel id associated with this value in HLO. Declared as `uint32`
	// even though `xla::HostCallbackArgInfo::channel_id` is `uint16_t` because
	// protobuf doesn't have a 16-bit integer type.
	ChannelId uint32 `protobuf:"varint,1,opt,name=channel_id,json=channelId,proto3" json:"channel_id,omitempty"`
	// The host shape for the value.
	Shape *data.ShapeProto `protobuf:"bytes,2,opt,name=shape,proto3" json:"shape,omitempty"`
	// contains filtered or unexported fields
}

func (*XlaHostCallbackProto_ArgInfo) Descriptor deprecated 

func (*XlaHostCallbackProto_ArgInfo) Descriptor() ([]byte, []int)

Deprecated: Use XlaHostCallbackProto_ArgInfo.ProtoReflect.Descriptor instead.

func (*XlaHostCallbackProto_ArgInfo) GetChannelId 

func (x *XlaHostCallbackProto_ArgInfo) GetChannelId() uint32

func (*XlaHostCallbackProto_ArgInfo) GetShape 

func (x *XlaHostCallbackProto_ArgInfo) GetShape() *data.ShapeProto

func (*XlaHostCallbackProto_ArgInfo) ProtoMessage 

func (*XlaHostCallbackProto_ArgInfo) ProtoMessage()

func (*XlaHostCallbackProto_ArgInfo) ProtoReflect 

func (x *XlaHostCallbackProto_ArgInfo) ProtoReflect() protoreflect.Message

func (*XlaHostCallbackProto_ArgInfo) Reset 

func (x *XlaHostCallbackProto_ArgInfo) Reset()

func (*XlaHostCallbackProto_ArgInfo) String 

func (x *XlaHostCallbackProto_ArgInfo) String() string

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.