Repository: ollama/ollama
File: ml/backend.go (Go)

Code before:
var backends = make(map[string]func(context.Context, *os.File, BackendParams) (Backend, error))

func RegisterBackend(name string, f func(context.Context, *os.File, BackendParams) (Backend, error)) {
	if _, ok := backends[name]; ok {
		panic("backend: backend already registered")
	}

	backends[name] = f
}

func NewBackend(ctx context.Context, f *os.File, params BackendParams) (Backend, error) {
	if backend, ok := backends["ggml"]; ok {
		return backend(ctx, f, params)
	}

	return nil, fmt.Errorf("unsupported backend")
}

type Context interface {
	Empty(dtype DType, shape ...int) Tensor
	Zeros(dtype DType, shape ...int) Tensor
	FromFloatSlice(s []float32, shape ...int) (Tensor, error)
	FromIntSlice(s []int32, shape ...int) (Tensor, error)
	Arange(start, stop, step float32, dtype DType) Tensor

	Forward(...Tensor) Context
	Compute(...Tensor)

	// Reserve is analogous to Compute but rather than executing a
	// graph, simply preallocates memory. Typically called with a
	// worst case graph to ensure all resources are available for
	// for future inference.
	Reserve() error

	MaxGraphNodes() int
	Close()

	// Input returns a context appropriate for creating tensors that are
	// inputs to the model (which includes things like output locations)
	Input() Context

	// Layer returns a context appropriate for creating intermediate tensors
	Layer(int) Context
}

type Tensor interface {
	Dim(n int) int
	Stride(n int) int

	Shape() []int
	DType() DType

Review comment (suggestion):
```suggestion
        // Arange returns a 1-D tensor with evenly spaced values from start to stop with the given step size; supports F32 (using ggml_arange) and I32 (using slice conversion) data types.
	Arange(start, stop, step float32, dtype DType) Tensor
```

Not super clear what this does if you're not familiar with other tensor libs

Code after:
var backends = make(map[string]func(context.Context, *os.File, BackendParams) (Backend, error))

func RegisterBackend(name string, f func(context.Context, *os.File, BackendParams) (Backend, error)) {
	if _, ok := backends[name]; ok {
		panic("backend: backend already registered")
	}

	backends[name] = f
}

func NewBackend(ctx context.Context, f *os.File, params BackendParams) (Backend, error) {
	if backend, ok := backends["ggml"]; ok {
		return backend(ctx, f, params)
	}

	return nil, fmt.Errorf("unsupported backend")
}

type Context interface {
	Empty(dtype DType, shape ...int) Tensor
	Zeros(dtype DType, shape ...int) Tensor
	FromFloatSlice(s []float32, shape ...int) (Tensor, error)
	FromIntSlice(s []int32, shape ...int) (Tensor, error)

	// Arange creates a 1D tensor with values within an interval (start, stop] increased by step.
	Arange(start, stop, step float32, dtype DType) Tensor

	Forward(...Tensor) Context
	Compute(...Tensor)

	// Reserve is analogous to Compute but rather than executing a
	// graph, simply preallocates memory. Typically called with a
	// worst case graph to ensure all resources are available for
	// for future inference.
	Reserve() error

	MaxGraphNodes() int
	Close()

	// Input returns a context appropriate for creating tensors that are
	// inputs to the model (which includes things like output locations)
	Input() Context

	// Layer returns a context appropriate for creating intermediate tensors
	Layer(int) Context
}

type Tensor interface {
	Dim(n int) int
	Stride(n int) int