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

↓ 77 callersMethodlen

examples/llama/var_store.rs:38

↓ 52 callersMethodas_ptr

src/wrappers/shape.rs:208

↓ 31 callersMethodclone

src/wrappers/xla_op.rs:51

↓ 26 callersMethodreshape

Reshape this node to a different set of dimension sizes, the number of element between the two different shapes has to match.

src/wrappers/xla_op.rs:145

↓ 24 callersFunctionhandle_status

(status: c_lib::status)

src/wrappers/mod.rs:294

↓ 24 callersMethodptr

src/wrappers/xla_builder.rs:22

↓ 23 callersMethodbuild

Generate a computation which root value is this node.

src/wrappers/xla_op.rs:648

↓ 21 callersMethodwrap

(&self, op: c_lib::xla_op)

src/wrappers/xla_op.rs:58

↓ 20 callersMethodcompile

Compile this computation for the specified client.

src/wrappers/mod.rs:324

↓ 20 callersMethodto_vec

Copy the values stored in the literal in a newly created vector. The data is flattened out for literals with more than one dimension.

src/wrappers/literal.rs:155

↓ 19 callersMethodparameter

Create an input node with the specified type and dimensions. A literal has to be passed for each of the parameter in the graph when calling the `execu

src/wrappers/xla_builder.rs:77

↓ 18 callersMethodto_literal_sync

Copy the buffer back to the host as a literal.

src/wrappers/pjrt_buffer.rs:27

↓ 16 callersMethodbuilder

src/wrappers/xla_op.rs:63

↓ 14 callersMethodwrap

(&self, op: c_lib::xla_op)

src/wrappers/xla_builder.rs:68

↓ 13 callersMethodc0

A shorter notation for `constant_r0`.

src/wrappers/xla_builder.rs:64

↓ 13 callersMethodget

↓ 12 callersFunctionget_and_remove

(key, transpose=False)

examples/llama/convert_checkpoint.py:14

↓ 12 callersMethodslice_in_dim1

A specialized version of `slice_in_dim` using a stride of one, so with all values with an index between `start_index` (inclusive) and `stop_index` (ex

src/wrappers/xla_op.rs:213

↓ 11 callersMethodarray_shape

src/wrappers/xla_op.rs:419

↓ 10 callersMethodconstant_r0

Create a node with a constant scalar value using the type of the element that is passed as argument.

src/wrappers/xla_builder.rs:58

↓ 10 callersMethodelement_type

src/wrappers/shape.rs:21

↓ 10 callersMethodprimitive_type

src/wrappers/mod.rs:121

↓ 10 callersMethodptr

src/wrappers/pjrt_client.rs:45

↓ 10 callersMethodswap_dims

Permute two dimensions, this is a specialized version of `transpose`.

src/wrappers/xla_op.rs:186

↓ 9 callersMethoddecompose_tuple

When the input is a tuple, return a vector of its elements. This replaces the original value by an empty tuple, no copy is performed.

src/wrappers/literal.rs:200

↓ 9 callersMethodto_string

↓ 8 callersMethodconvert

Create a new node by casting the elements of the original node to a new primitive type.

src/wrappers/xla_op.rs:302

↓ 8 callersMethoddims

src/wrappers/xla_op.rs:423

↓ 8 callersMethodprimitive_type

The kind of elements that are computed by this operand.

src/wrappers/xla_op.rs:401

↓ 8 callersMethodreduce_sum

A node that computes the sum across the specified dimensions, e.g. if all the dimensions are passed as an argument the result is a scalar with the sum

src/wrappers/xla_op.rs:529

↓ 8 callersMethodtake

( &mut self, s: &str, expected_type: ElementType, expected_dims: &[usize],

examples/nanogpt/var_store.rs:24

↓ 7 callersFunctionc_ptr_to_string

(ptr: *const std::ffi::c_char)

src/wrappers/mod.rs:24

↓ 7 callersMethoddims

src/wrappers/shape.rs:40

↓ 7 callersMethodlen

examples/nanogpt/var_store.rs:20

↓ 7 callersMethodnormalize_index

(&self, index: i64)

src/wrappers/xla_op.rs:325

↓ 6 callersMethodbroadcast

Add some broadcasting dimensions at the beginning of the current node shape.

src/wrappers/xla_op.rs:151

↓ 6 callersMethodshape

src/wrappers/shape.rs:181

↓ 5 callersMethodconstant_literal

Create a node with a constant value defined by the specified literal.

src/wrappers/xla_builder.rs:51

↓ 5 callersMethoddevices

A list of devices attached to this client.

src/wrappers/pjrt_client.rs:84

↓ 5 callersMethodrank

The number of dimensions for this node.

src/wrappers/xla_op.rs:411

↓ 5 callersMethodty

src/wrappers/shape.rs:25

↓ 5 callersMethodvar

(&mut self, s: &str, dims: &[usize])

examples/llama/var_store.rs:58

↓ 4 callersMethodconstant_r1

A one dimension constant node based on some slice stored on the host.

src/wrappers/xla_builder.rs:123

↓ 4 callersMethoddimensions_size

A node that contains the size of the dimension with the target index as a `S32` scalar value.

src/wrappers/xla_op.rs:340

↓ 4 callersMethoddot_general

General dot multiplication between two nodes, specifying the dimensions that get contracted as well as the batch dimensions.

src/wrappers/xla_op.rs:440

↓ 4 callersMethodelement_count

The number of elements stored in the literal.

src/wrappers/literal.rs:68

↓ 4 callersFunctionmake_shape_internal

xla_rs/xla_rs.cc:213

↓ 4 callersMethodmatmul

Matrix multiplication, this is a specialized version of `dot_general` to be used for matrix-matrix or matrix-vector multiplications.

src/wrappers/xla_op.rs:594

↓ 4 callersFunctionparse

↓ 4 callersMethodreduce_mean

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

src/wrappers/xla_op.rs:540

↓ 4 callersMethodsoftmax

(&self, dim: i64)

src/wrappers/xla_op.rs:573

↓ 4 callersMethodto_string

src/wrappers/pjrt_device.rs:25

↓ 4 callersMethodtuple

Create a new tuple shape.

src/wrappers/shape.rs:73

↓ 3 callersMethodelement_count

The number of elements stored in arrays that use this shape, this is the product of sizes across each dimension.

src/wrappers/shape.rs:36

↓ 3 callersMethodelement_size_in_bytes

The size for this element type in bytes.

src/wrappers/mod.rs:101

↓ 3 callersMethodexecute

( &self, args: &[L], )

src/wrappers/pjrt_loaded_executable.rs:42

↓ 3 callersMethodget_tuple_element

Index into tuples.

src/wrappers/xla_op.rs:253

↓ 3 callersMethodone

A scalar node with the one value for the associated type.

src/wrappers/xla_builder.rs:140

↓ 3 callersMethodparameter_s

(&self, parameter_number: i64, shape: &Shape, name: &str)

src/wrappers/xla_builder.rs:107

↓ 3 callersFunctionread_header

(reader: &mut R)

↓ 3 callersMethodreduce

Create a node by folding a computation acress some target dimensions. If `keep_dims` is `true`, the resulting node has a dimension of size one for the

src/wrappers/xla_op.rs:350

↓ 3 callersFunctiontr

examples/llama/convert_checkpoint.py:8

↓ 3 callersMethodtranspose

Permute the dimension with the specified indexes.

src/wrappers/xla_op.rs:180

↓ 3 callersMethodty

The element type used by element stored in this literal, shortcut for `element_type`.

src/wrappers/literal.rs:87

↓ 3 callersMethodzero

A scalar node with the zero value for the associated type.

src/wrappers/xla_builder.rs:134

↓ 2 callersMethodaddressable_device_count

The number of devices that this client can use.

src/wrappers/pjrt_client.rs:63

↓ 2 callersMethodapply_rotary_emb

(&self, x: &XlaOp, freqs_cis: &XlaOp)

examples/llama/main.rs:218

↓ 2 callersMethodbuffer_from_host_raw_bytes

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:137

↓ 2 callersMethodclient

The client that owns this buffer.

src/wrappers/pjrt_buffer.rs:13

↓ 2 callersMethodconcat_in_dim

Concat multiple nodes (together with the `self` node) along the target dimension.

src/wrappers/xla_op.rs:241

↓ 2 callersMethodconstant_r1c

(&self, f: T, len: usize)

src/wrappers/xla_builder.rs:117

↓ 2 callersMethodcopy_raw_to

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

src/wrappers/literal.rs:112

↓ 2 callersMethoddevice_count

The number of devices that this client has detected, e.g. the number of GPUs.

src/wrappers/pjrt_client.rs:58

↓ 2 callersMethodget_current_status

This returns `Ok(())` if the graph creation has not generated any error so far. Otherwise the current status is returned.

src/wrappers/xla_builder.rs:44

↓ 2 callersMethodget_dimensions_size

The number of dimensions (a.k.a the rank) associated with this op.

src/wrappers/xla_builder.rs:235

↓ 2 callersMethodget_shape

The shape associated with this op.

src/wrappers/xla_builder.rs:209

↓ 2 callersMethodlast_dim

src/wrappers/shape.rs:48

↓ 2 callersMethodname

The computation name.

src/wrappers/mod.rs:316

↓ 2 callersMethodon_device_shape

Retrieve the shape used by this buffer.

src/wrappers/pjrt_buffer.rs:35

↓ 2 callersFunctionpjrt_buffer_from_host_buffer

xla_rs/xla_rs.cc:94

↓ 2 callersMethodprimitive_type

The primitive type used by element stored in this literal.

src/wrappers/literal.rs:73

↓ 2 callersMethodprocess_execute_outputs

( &self, outputs: *mut *mut c_lib::pjrt_buffer, )

src/wrappers/pjrt_loaded_executable.rs:15

↓ 2 callersMethodselect

Select values from the original tensor to be values from `on_true` if the associated value in `self` is true, and the values from `on_false` otherwise

src/wrappers/xla_op.rs:266

↓ 2 callersFunctionset_tf_min_log_level

(log_level: TfLogLevel)

↓ 2 callersMethodshape

The [`Shape`] of the literal, this contains information about the dimensions of the underlying array, as well as the primitive type of the array's ele

src/wrappers/literal.rs:99

↓ 2 callersMethodsize_bytes

The literal size in bytes, this is the same as `element_count` multiplied by `element_size_in_bytes`.

src/wrappers/literal.rs:93

↓ 2 callersMethodslice_in_dim

Create a node that has a partial view on the data of the original node. Indexes on the target dimension `dim` are restricted to the values between `st

src/wrappers/xla_op.rs:199

↓ 2 callersMethodvar_

( &mut self, s: &str, ty: ElementType, dims: &[usize], is_arg: bool,

examples/llama/var_store.rs:42

↓ 2 callersMethodwrite

(&self, f: &mut T)

↓ 1 callersMethodaddressable_devices

A list of devices that can be used by this client.

src/wrappers/pjrt_client.rs:92

↓ 1 callersMethodarg

(&mut self, s: &str, ty: ElementType, dims: &[usize])

examples/llama/var_store.rs:63

↓ 1 callersMethodarg_indexes

examples/llama/var_store.rs:98

↓ 1 callersMethodas_env_variable_str

↓ 1 callersMethodbpe

(&self, s: &str)

examples/llama/sentencepiece.rs:88

↓ 1 callersMethodbpe

(&self, token: &[u8])

examples/nanogpt/tokenizer.rs:333

↓ 1 callersMethodbuffer_from_host_buffer

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:104

↓ 1 callersFunctionbuild

xla_rs/xla_rs.cc:864

↓ 1 callersMethodc1

Shorthand function for `constant_r1`.

src/wrappers/xla_builder.rs:129

↓ 1 callersMethodc_shape

src/wrappers/shape.rs:101

↓ 1 callersFunctioncompile

xla_rs/xla_rs.cc:871

next →1–100 of 409, ranked by callers