Functions — github.com/LaurentMazare/xla-rs

CodeHub MCPcopy Create free account

hub / github.com/LaurentMazare/xla-rs / functions

Functions409 in github.com/LaurentMazare/xla-rs

⨍ Functions 409◇ Types & classes 59

↓ 1 callersFunctionpjrt_client_addressable_device_count

xla_rs/xla_rs.cc:63

↓ 1 callersFunctionpjrt_client_addressable_devices

xla_rs/xla_rs.cc:74

↓ 1 callersFunctionpjrt_client_device_count

xla_rs/xla_rs.cc:61

↓ 1 callersFunctionpjrt_client_devices

xla_rs/xla_rs.cc:67

↓ 1 callersFunctionpjrt_client_free

xla_rs/xla_rs.cc:90

↓ 1 callersFunctionpjrt_client_platform_name

xla_rs/xla_rs.cc:81

↓ 1 callersFunctionpjrt_client_platform_version

xla_rs/xla_rs.cc:86

↓ 1 callersFunctionpjrt_cpu_client_create

xla_rs/xla_rs.cc:35

↓ 1 callersFunctionpjrt_device_debug_string

xla_rs/xla_rs.cc:166

↓ 1 callersFunctionpjrt_device_id

xla_rs/xla_rs.cc:144

↓ 1 callersFunctionpjrt_device_kind

xla_rs/xla_rs.cc:162

↓ 1 callersFunctionpjrt_device_local_hardware_id

xla_rs/xla_rs.cc:148

↓ 1 callersFunctionpjrt_device_process_index

xla_rs/xla_rs.cc:146

↓ 1 callersFunctionpjrt_device_to_string

xla_rs/xla_rs.cc:170

↓ 1 callersFunctionpjrt_device_transfer_from_outfeed

xla_rs/xla_rs.cc:157

↓ 1 callersFunctionpjrt_device_transfer_to_infeed

xla_rs/xla_rs.cc:152

↓ 1 callersFunctionpjrt_gpu_client_create

xla_rs/xla_rs.cc:41

↓ 1 callersFunctionpjrt_loaded_executable_free

xla_rs/xla_rs.cc:92

↓ 1 callersMethodplatform_name

The name of the platform.

src/wrappers/pjrt_client.rs:68

↓ 1 callersMethodplatform_version

The version of the platform.

src/wrappers/pjrt_client.rs:76

↓ 1 callersFunctionprecompute_freqs_cis

(config: &Config, builder: &XlaBuilder)

examples/llama/main.rs:328

↓ 1 callersMethodprimitive_type

The stored primitive type.

src/wrappers/shape.rs:30

↓ 1 callersMethodprocess_index

src/wrappers/pjrt_device.rs:16

↓ 1 callersMethodproto

Get the HloModuleProto for the computation.

src/wrappers/mod.rs:329

↓ 1 callersMethodreduce_max

A node that computes the maximum value across the specified dimensions.

src/wrappers/xla_op.rs:552

↓ 1 callersFunctionsample

(exe: &PjRtLoadedExecutable, tokenizer: &Tokenizer, cnt: usize)

examples/nanogpt/main.rs:262

↓ 1 callersFunctionshape_dimensions

xla_rs/xla_rs.cc:832

↓ 1 callersFunctionshape_dimensions_size

xla_rs/xla_rs.cc:828

↓ 1 callersFunctionshape_element_type

xla_rs/xla_rs.cc:830

↓ 1 callersFunctionshape_free

xla_rs/xla_rs.cc:834

↓ 1 callersFunctionshape_tuple_shapes

xla_rs/xla_rs.cc:824

↓ 1 callersFunctionshape_tuple_shapes_size

xla_rs/xla_rs.cc:822

↓ 1 callersMethodsilu

SiLU activation function. This computes the element-wise SiLU activation, x.sigmoid(x).

src/wrappers/xla_op.rs:124

↓ 1 callersMethodsqueeze

Squeeze the dimension as the target index, i.e. if this dimension has size one remove it for the generated node. The target dimension index can be spe

src/wrappers/xla_op.rs:228

↓ 1 callersFunctionstatus_error_message

xla_rs/xla_rs.cc:1035

↓ 1 callersFunctionstatus_free

xla_rs/xla_rs.cc:1027

↓ 1 callersMethodto_tuple1

src/wrappers/literal.rs:216

↓ 1 callersFunctionxla_builder_create

xla_rs/xla_rs.cc:174

↓ 1 callersFunctionxla_builder_free

xla_rs/xla_rs.cc:178

↓ 1 callersFunctionxla_computation_free

xla_rs/xla_rs.cc:1033

↓ 1 callersFunctionxla_computation_from_hlo_module_proto

xla_rs/xla_rs.cc:1076

↓ 1 callersFunctionxla_computation_name

xla_rs/xla_rs.cc:1029

↓ 1 callersFunctionxla_computation_proto

xla_rs/xla_rs.cc:1083

↓ 1 callersFunctionxla_op_free

xla_rs/xla_rs.cc:820

tests/basic_tests.rs:4

Create a new array shape.

src/wrappers/shape.rs:63

Methodarray_shape

src/wrappers/literal.rs:106

Methodarray_with_type

Create a new array shape.

src/wrappers/shape.rs:68

A new node containing only values for index `index_in_dim` on the dimension `dim_index`. The target dimension is squeezed so the resulting node has on

src/wrappers/xla_op.rs:220

Methodbroadcast_in_dim

Add some broadcasting dimensions at arbitrary positions. See the [semantics](https://www.tensorflow.org/xla/operation_semantics#broadcastindim).

src/wrappers/xla_op.rs:159

Methodbuffer_from_host_literal

Transfer some data from the host to a `PjRtBuffer` stored on the target device. If the device is not specified, the default device is used. The source

src/wrappers/pjrt_client.rs:170

Build a computation from the specified root node. This can only be called once.

src/wrappers/xla_builder.rs:27

Clamp the values in the original node to be between `min` and `max`.

src/wrappers/xla_op.rs:259

The client that owns this executable.

src/wrappers/pjrt_loaded_executable.rs:11

src/wrappers/literal.rs:10

Collapse the dimensions of this node into a single dimension, [xla documentation](https://www.tensorflow.org/xla/operation_semantics#collapse).

src/wrappers/xla_op.rs:174

Compile a computation for this device, and return the executable.

src/wrappers/pjrt_client.rs:50

Methodconditional

Execute `true_comp` if `self` is true, `false_comp` if `self` is false, and return the result. `self` has to be a scalar of type `PRED`. `true_op` is

src/wrappers/xla_op.rs:380

Methodconfig_13b

examples/llama/main.rs:92

Methodconfig_30b

examples/llama/main.rs:96

Methodconfig_65b

examples/llama/main.rs:100

Methodconfig_7b

examples/llama/main.rs:88

Create a new literal containing the data from the original literal casted to a new primitive type. The dimensions of the resulting literal are the sam

src/wrappers/literal.rs:191

Methodcopy_raw_from

Copy data from a slice to the literal. This returns an error if the primitive type used by the literal is not `T` or if number of elements in the slic

src/wrappers/literal.rs:134

Methodcopy_raw_to_host_sync

Copy the data stored in a buffer to host memory in a blocking way.

src/wrappers/pjrt_buffer.rs:42

Methodcopy_to_device

Copy the buffer to a different device.

src/wrappers/pjrt_buffer.rs:18

A CPU client, this can run computations on multiple CPUs at the same time.

src/wrappers/pjrt_client.rs:16

Methodcreate_from_shape

Create an unitialized literal based on some primitive type and some dimensions.

src/wrappers/literal.rs:18

Methodcreate_from_shape_and_untyped_data

Create an unitialized literal based on some primitive type, some dimensions, and some data. The data is untyped, i.e. it is a sequence of bytes repres

src/wrappers/literal.rs:27

Methoddebug_string

src/wrappers/pjrt_device.rs:39

(&self, tokens: &[usize])

examples/llama/sentencepiece.rs:146

examples/nanogpt/main.rs:192

examples/llama/var_store.rs:76

examples/nanogpt/var_store.rs:62

src/wrappers/mod.rs:336

src/wrappers/literal.rs:281

src/wrappers/pjrt_loaded_executable.rs:68

src/wrappers/pjrt_client.rs:186

src/wrappers/pjrt_buffer.rs:70

src/wrappers/xla_builder.rs:251

src/wrappers/xla_op.rs:654

src/wrappers/shape.rs:214

A node that applies the specified Einstein summation formula to this node.

src/wrappers/xla_op.rs:129

A node that applies the specified Einstein summation formula to this node and the other argument node.

src/wrappers/xla_op.rs:137

Methodelement_type

src/wrappers/mod.rs:56

Methodfirst_dim

src/wrappers/shape.rs:44

src/wrappers/literal.rs:269

Methodfrom_file

examples/llama/sentencepiece.rs:15

Methodfrom_proto

(proto: &HloModuleProto)

src/wrappers/mod.rs:310

Methodfrom_proto_file

Read a HLO module from a proto file, either in binary or pbtxt format.

src/wrappers/mod.rs:354

(ptr: c_lib::shape)

src/wrappers/shape.rs:177

Methodfrom_ptr_rec

(ptr: c_lib::shape)

src/wrappers/shape.rs:182

Methodfrom_raw_bytes

( _: &Self::Context, ty: ElementType, dims: &[usize], bytes: &[u8], )

Methodfrom_text_file

Read a HLO module from a text file.

src/wrappers/mod.rs:345

Methodget_first_element

Get the first element from a literal. This returns an error if type `T` is not the primitive type that the literal uses.

src/wrappers/literal.rs:55

Methodget_pairs

(word: &[Vec<u8>])

examples/llama/sentencepiece.rs:78

Methodget_pairs

(word: &[Vec<u8>])

examples/nanogpt/tokenizer.rs:323

A GPU client, the memory requirements are limited by the specified `memory_fraction` and this memory can either be allocated dynamically or pre-alloca

src/wrappers/pjrt_client.rs:26

Read a single value from the implicit streaming interface of the device.

src/wrappers/xla_builder.rs:99

Methodinvalid_argument_error

An error node, using the 'invalid argument error' error type.

src/wrappers/xla_builder.rs:193

← previous next →201–300 of 409, ranked by callers