embeddings static method

Tensor embeddings(
  1. List<int> idx,
  2. Tensor wte,
  3. Tensor wpe
)

Implementation

static Tensor embeddings(List<int> idx, Tensor wte, Tensor wpe) {
  final T = idx.length;
  final D = wte.shape[1];

  // Allocate memory for the indices that C can read
  final ptr = calloc<ffi.Int32>(T);
  for (int i = 0; i < T; i++) {
    ptr[i] = idx[i];
  }

  final handle = engine.embeddingForward(ptr, wte._handle, wpe._handle, T, D);

  // We can free the host pointer immediately because C++ copied it to d_indices
  calloc.free(ptr);

  return Tensor._raw(handle, [T, D]);
}