@hoff97/tensor-js
Version:
PyTorch like deep learning inferrence library
36 lines (35 loc) • 1.72 kB
TypeScript
import { DTypeGpu, GPUTensorConstructor, GPUTensorI } from '../../../tensor/gpu/interface';
import { GPUMemoryAllocator } from '../../../tensor/gpu/memory';
import { GPUTensor } from '../../../tensor/gpu/tensor';
import { Dispatcher } from '../../gpu/dispatcher';
import { Input, Operation } from '../../gpu/operation';
export interface RepeatIndexInfo {
shapeA?: readonly number[];
widthA?: number;
heightA?: number;
shapeOutput?: readonly number[];
widthOutput?: number;
heightOutput?: number;
sparseShape?: readonly number[];
repeatStrides?: readonly number[];
}
export interface RepeatIndexInput {
A: GPUTensorI;
repeats: readonly number[];
shape: readonly number[];
repeatsProd: number;
}
export declare class RepeatIndexOperation<GPUTensor extends GPUTensorI> extends Operation<GPUTensor, RepeatIndexInfo, RepeatIndexInput> {
constructor(tensorConstructor: GPUTensorConstructor<GPUTensor>, dtype: DTypeGpu, allocator?: GPUMemoryAllocator);
getVariables(): string;
getUniformAttrs(): Input[];
getFragmentShader(info: RepeatIndexInfo): string;
getTextureNames(): string[];
calc(input: RepeatIndexInput): GPUTensor;
getOutputShape(input: RepeatIndexInput): readonly number[];
compile(info: RepeatIndexInfo): void;
getCompilationInfo(input: RepeatIndexInput): RepeatIndexInfo;
getInputInfoString(input: RepeatIndexInput): string;
}
export declare const defaultRepeatIndexD: Dispatcher<GPUTensorI, unknown, unknown, RepeatIndexOperation<GPUTensor<DTypeGpu>>>;
export declare function repeatIndexGPU(indices: GPUTensor<'uint32'>, repeats: readonly number[], shape: readonly number[], repeatsProd: number): GPUTensor<'uint32'>;