server: add tests and fix isHuggingFaceURL edge case

- Add comprehensive tests for isHuggingFaceURL and getNumDownloadParts
- Fix bug where domains ending in huggingface.co (like nothuggingface.co)
  would incorrectly match as HuggingFace URLs
- Improve code comments with more detailed documentation

CMakeLists.txt

@@ -190,7 +190,7 @@ if(MLX_ENGINE)
     install(TARGETS mlx mlxc
         RUNTIME_DEPENDENCIES
             DIRECTORIES ${CUDAToolkit_BIN_DIR} ${CUDAToolkit_BIN_DIR}/x64 ${CUDAToolkit_LIBRARY_DIR}
-            PRE_INCLUDE_REGEXES cublas cublasLt cudart nvrtc nvrtc-builtins cudnn nccl openblas gfortran
+            PRE_INCLUDE_REGEXES cublas cublasLt cudart nvrtc cudnn nccl
             PRE_EXCLUDE_REGEXES ".*"
         RUNTIME DESTINATION ${OLLAMA_INSTALL_DIR} COMPONENT MLX
         LIBRARY DESTINATION ${OLLAMA_INSTALL_DIR} COMPONENT MLX

Dockerfile

@@ -32,7 +32,7 @@ ENV PATH=/${VULKANVERSION}/x86_64/bin:$PATH
 FROM --platform=linux/arm64 almalinux:8 AS base-arm64
 # install epel-release for ccache
 RUN yum install -y yum-utils epel-release \
-    && dnf install -y clang ccache git \
+    && dnf install -y clang ccache \
     && yum-config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/sbsa/cuda-rhel8.repo
 ENV CC=clang CXX=clang++
 
@@ -149,7 +149,6 @@ COPY CMakeLists.txt CMakePresets.json .
 COPY ml/backend/ggml/ggml ml/backend/ggml/ggml
 COPY x/ml/backend/mlx x/ml/backend/mlx
 COPY go.mod go.sum .
-COPY MLX_VERSION .
 RUN curl -fsSL https://golang.org/dl/go$(awk '/^go/ { print $2 }' go.mod).linux-$(case $(uname -m) in x86_64) echo amd64 ;; aarch64) echo arm64 ;; esac).tar.gz | tar xz -C /usr/local
 ENV PATH=/usr/local/go/bin:$PATH
 RUN go mod download
@@ -157,6 +156,14 @@ RUN --mount=type=cache,target=/root/.ccache \
     cmake --preset 'MLX CUDA 13' -DBLAS_INCLUDE_DIRS=/usr/include/openblas -DLAPACK_INCLUDE_DIRS=/usr/include/openblas \
         && cmake --build --parallel ${PARALLEL} --preset 'MLX CUDA 13' \
         && cmake --install build --component MLX --strip --parallel ${PARALLEL}
+COPY . .
+ARG GOFLAGS="'-ldflags=-w -s'"
+ENV CGO_ENABLED=1
+ARG CGO_CFLAGS
+ARG CGO_CXXFLAGS
+RUN mkdir -p dist/bin
+RUN --mount=type=cache,target=/root/.cache/go-build \
+    go build -tags mlx -trimpath -buildmode=pie -o dist/bin/ollama-mlx .
 
 FROM base AS build
 WORKDIR /go/src/github.com/ollama/ollama
@@ -165,14 +172,12 @@ RUN curl -fsSL https://golang.org/dl/go$(awk '/^go/ { print $2 }' go.mod).linux-
 ENV PATH=/usr/local/go/bin:$PATH
 RUN go mod download
 COPY . .
-# Clone mlx-c headers for CGO (version from MLX_VERSION file)
-RUN git clone --depth 1 --branch "$(cat MLX_VERSION)" https://github.com/ml-explore/mlx-c.git build/_deps/mlx-c-src
 ARG GOFLAGS="'-ldflags=-w -s'"
 ENV CGO_ENABLED=1
-ENV CGO_CFLAGS="-I/go/src/github.com/ollama/ollama/build/_deps/mlx-c-src"
+ARG CGO_CFLAGS
 ARG CGO_CXXFLAGS
 RUN --mount=type=cache,target=/root/.cache/go-build \
-    go build -tags mlx -trimpath -buildmode=pie -o /bin/ollama .
+    go build -trimpath -buildmode=pie -o /bin/ollama .
 
 FROM --platform=linux/amd64 scratch AS amd64
 # COPY --from=cuda-11 dist/lib/ollama/ /lib/ollama/
@@ -180,6 +185,7 @@ COPY --from=cuda-12 dist/lib/ollama /lib/ollama/
 COPY --from=cuda-13 dist/lib/ollama /lib/ollama/
 COPY --from=vulkan  dist/lib/ollama  /lib/ollama/
 COPY --from=mlx     /go/src/github.com/ollama/ollama/dist/lib/ollama /lib/ollama/
+COPY --from=mlx     /go/src/github.com/ollama/ollama/dist/bin/ /bin/
 
 FROM --platform=linux/arm64 scratch AS arm64
 # COPY --from=cuda-11 dist/lib/ollama/ /lib/ollama/

MLX_VERSION

@@ -1 +0,0 @@
-v0.4.1

README.md

@@ -48,7 +48,7 @@ ollama run gemma3
 
 ## Model library
 
-Ollama supports a list of models available on [ollama.com/library](https://ollama.com/library "ollama model library")
+Ollama supports a list of models available on [ollama.com/library](https://ollama.com/library 'ollama model library')
 
 Here are some example models that can be downloaded:
 
@@ -79,7 +79,7 @@ Here are some example models that can be downloaded:
 | Code Llama         | 7B         | 3.8GB | `ollama run codellama`           |
 | Llama 2 Uncensored | 7B         | 3.8GB | `ollama run llama2-uncensored`   |
 | LLaVA              | 7B         | 4.5GB | `ollama run llava`               |
-| Granite-3.3        | 8B         | 4.9GB | `ollama run granite3.3`          |
+| Granite-3.3         | 8B         | 4.9GB | `ollama run granite3.3`          |
 
 > [!NOTE]
 > You should have at least 8 GB of RAM available to run the 7B models, 16 GB to run the 13B models, and 32 GB to run the 33B models.
@@ -260,38 +260,6 @@ Finally, in a separate shell, run a model:
 ./ollama run llama3.2
 ```
 
-## Building with MLX (experimental)
-
-First build the MLX libraries:
-
-```shell
-cmake --preset MLX
-cmake --build --preset MLX --parallel
-cmake --install build --component MLX
-```
-
-When building with the `-tags mlx` flag, the main `ollama` binary includes MLX support for experimental features like image generation:
-
-```shell
-go build -tags mlx .
-```
-
-Finally, start the server:
-
-```
-./ollama serve
-```
-
-### Building MLX with CUDA
-
-When building with CUDA, use the preset "MLX CUDA 13" or "MLX CUDA 12" to enable CUDA with default architectures:
-
-```shell
-cmake --preset 'MLX CUDA 13'
-cmake --build --preset 'MLX CUDA 13' --parallel
-cmake --install build --component MLX
-```
-
 ## REST API
 
 Ollama has a REST API for running and managing models.
@@ -322,7 +290,6 @@ See the [API documentation](./docs/api.md) for all endpoints.
 
 ### Web & Desktop
 
-- [Onyx](https://github.com/onyx-dot-app/onyx)
 - [Open WebUI](https://github.com/open-webui/open-webui)
 - [SwiftChat (macOS with ReactNative)](https://github.com/aws-samples/swift-chat)
 - [Enchanted (macOS native)](https://github.com/AugustDev/enchanted)
@@ -454,7 +421,7 @@ See the [API documentation](./docs/api.md) for all endpoints.
 - [AppFlowy](https://github.com/AppFlowy-IO/AppFlowy) (AI collaborative workspace with Ollama, cross-platform and self-hostable)
 - [Lumina](https://github.com/cushydigit/lumina.git) (A lightweight, minimal React.js frontend for interacting with Ollama servers)
 - [Tiny Notepad](https://pypi.org/project/tiny-notepad) (A lightweight, notepad-like interface to chat with ollama available on PyPI)
-- [macLlama (macOS native)](https://github.com/hellotunamayo/macLlama) (A native macOS GUI application for interacting with Ollama models, featuring a chat interface.)
+- [macLlama (macOS native)](https://github.com/hellotunamayo/macLlama) (A native macOS GUI application for interacting with Ollama models, featuring a chat interface.) 
 - [GPTranslate](https://github.com/philberndt/GPTranslate) (A fast and lightweight, AI powered desktop translation application written with Rust and Tauri. Features real-time translation with OpenAI/Azure/Ollama.)
 - [ollama launcher](https://github.com/NGC13009/ollama-launcher) (A launcher for Ollama, aiming to provide users with convenient functions such as ollama server launching, management, or configuration.)
 - [ai-hub](https://github.com/Aj-Seven/ai-hub) (AI Hub supports multiple models via API keys and Chat support via Ollama API.)
@@ -526,7 +493,7 @@ See the [API documentation](./docs/api.md) for all endpoints.
 ### Database
 
 - [pgai](https://github.com/timescale/pgai) - PostgreSQL as a vector database (Create and search embeddings from Ollama models using pgvector)
-  - [Get started guide](https://github.com/timescale/pgai/blob/main/docs/vectorizer-quick-start.md)
+   - [Get started guide](https://github.com/timescale/pgai/blob/main/docs/vectorizer-quick-start.md)
 - [MindsDB](https://github.com/mindsdb/mindsdb/blob/staging/mindsdb/integrations/handlers/ollama_handler/README.md) (Connects Ollama models with nearly 200 data platforms and apps)
 - [chromem-go](https://github.com/philippgille/chromem-go/blob/v0.5.0/embed_ollama.go) with [example](https://github.com/philippgille/chromem-go/tree/v0.5.0/examples/rag-wikipedia-ollama)
 - [Kangaroo](https://github.com/dbkangaroo/kangaroo) (AI-powered SQL client and admin tool for popular databases)
@@ -669,7 +636,6 @@ See the [API documentation](./docs/api.md) for all endpoints.
 - [llama.cpp](https://github.com/ggml-org/llama.cpp) project founded by Georgi Gerganov.
 
 ### Observability
-
 - [Opik](https://www.comet.com/docs/opik/cookbook/ollama) is an open-source platform to debug, evaluate, and monitor your LLM applications, RAG systems, and agentic workflows with comprehensive tracing, automated evaluations, and production-ready dashboards. Opik supports native integration to Ollama.
 - [Lunary](https://lunary.ai/docs/integrations/ollama) is the leading open-source LLM observability platform. It provides a variety of enterprise-grade features such as real-time analytics, prompt templates management, PII masking, and comprehensive agent tracing.
 - [OpenLIT](https://github.com/openlit/openlit) is an OpenTelemetry-native tool for monitoring Ollama Applications & GPUs using traces and metrics.
@@ -678,5 +644,4 @@ See the [API documentation](./docs/api.md) for all endpoints.
 - [MLflow Tracing](https://mlflow.org/docs/latest/llms/tracing/index.html#automatic-tracing) is an open source LLM observability tool with a convenient API to log and visualize traces, making it easy to debug and evaluate GenAI applications.
 
 ### Security
-
 - [Ollama Fortress](https://github.com/ParisNeo/ollama_proxy_server)

api/types.go

@@ -127,20 +127,6 @@ type GenerateRequest struct {
 	// each with an associated log probability. Only applies when Logprobs is true.
 	// Valid values are 0-20. Default is 0 (only return the selected token's logprob).
 	TopLogprobs int `json:"top_logprobs,omitempty"`
-
-	// Experimental: Image generation fields (may change or be removed)
-
-	// Width is the width of the generated image in pixels.
-	// Only used for image generation models.
-	Width int32 `json:"width,omitempty"`
-
-	// Height is the height of the generated image in pixels.
-	// Only used for image generation models.
-	Height int32 `json:"height,omitempty"`
-
-	// Steps is the number of diffusion steps for image generation.
-	// Only used for image generation models.
-	Steps int32 `json:"steps,omitempty"`
 }
 
 // ChatRequest describes a request sent by [Client.Chat].
@@ -874,20 +860,6 @@ type GenerateResponse struct {
 	// Logprobs contains log probability information for the generated tokens,
 	// if requested via the Logprobs parameter.
 	Logprobs []Logprob `json:"logprobs,omitempty"`
-
-	// Experimental: Image generation fields (may change or be removed)
-
-	// Image contains a base64-encoded generated image.
-	// Only present for image generation models.
-	Image string `json:"image,omitempty"`
-
-	// Completed is the number of completed steps in image generation.
-	// Only present for image generation models during streaming.
-	Completed int64 `json:"completed,omitempty"`
-
-	// Total is the total number of steps for image generation.
-	// Only present for image generation models during streaming.
-	Total int64 `json:"total,omitempty"`
 }
 
 // ModelDetails provides details about a model.

app/cmd/app/app_darwin.m

@@ -14,7 +14,6 @@ extern NSString *SystemWidePath;
 @interface AppDelegate () <NSWindowDelegate, WKNavigationDelegate, WKUIDelegate>
 @property(strong, nonatomic) NSStatusItem *statusItem;
 @property(assign, nonatomic) BOOL updateAvailable;
-@property(assign, nonatomic) BOOL systemShutdownInProgress;
 @end
 
 @implementation AppDelegate
@@ -41,13 +40,6 @@ bool firstTimeRun,startHidden; // Set in run before initialization
 }
 
 - (void)applicationDidFinishLaunching:(NSNotification *)aNotification {
-    // Register for system shutdown/restart notification so we can allow termination
-    [[[NSWorkspace sharedWorkspace] notificationCenter]
-        addObserver:self
-           selector:@selector(systemWillPowerOff:)
-               name:NSWorkspaceWillPowerOffNotification
-             object:nil];
-
     // if we're in development mode, set the app icon
     NSString *bundlePath = [[NSBundle mainBundle] bundlePath];
     if (![bundlePath hasSuffix:@".app"]) {
@@ -286,18 +278,7 @@ bool firstTimeRun,startHidden; // Set in run before initialization
     [NSApp activateIgnoringOtherApps:YES];
 }
 
-- (void)systemWillPowerOff:(NSNotification *)notification {
-    // Set flag so applicationShouldTerminate: knows to allow termination.
-    // The system will call applicationShouldTerminate: after posting this notification.
-    self.systemShutdownInProgress = YES;
-}
-
 - (NSApplicationTerminateReply)applicationShouldTerminate:(NSApplication *)sender {
-    // Allow termination if the system is shutting down or restarting
-    if (self.systemShutdownInProgress) {
-        return NSTerminateNow;
-    }
-    // Otherwise just hide the app (for Cmd+Q, close button, etc.)
     [NSApp hide:nil];
     [NSApp setActivationPolicy:NSApplicationActivationPolicyAccessory];
     return NSTerminateCancel;

cmd/cmd.go

@@ -46,9 +46,8 @@ import (
 	"github.com/ollama/ollama/types/syncmap"
 	"github.com/ollama/ollama/version"
 	xcmd "github.com/ollama/ollama/x/cmd"
-	"github.com/ollama/ollama/x/create"
-	xcreateclient "github.com/ollama/ollama/x/create/client"
 	"github.com/ollama/ollama/x/imagegen"
+	imagegenclient "github.com/ollama/ollama/x/imagegen/client"
 )
 
 const ConnectInstructions = "To sign in, navigate to:\n    %s\n\n"
@@ -94,87 +93,15 @@ func CreateHandler(cmd *cobra.Command, args []string) error {
 	p := progress.NewProgress(os.Stderr)
 	defer p.Stop()
 
-	// Validate model name early to fail fast
-	modelName := args[0]
-	name := model.ParseName(modelName)
-	if !name.IsValid() {
-		return fmt.Errorf("invalid model name: %s", modelName)
-	}
-
-	// Check for --experimental flag for safetensors model creation
-	experimental, _ := cmd.Flags().GetBool("experimental")
-	if experimental {
-		// Get Modelfile content - either from -f flag or default to "FROM ."
-		var reader io.Reader
-		filename, err := getModelfileName(cmd)
-		if os.IsNotExist(err) || filename == "" {
-			// No Modelfile specified or found - use default
-			reader = strings.NewReader("FROM .\n")
-		} else if err != nil {
-			return err
-		} else {
-			f, err := os.Open(filename)
-			if err != nil {
-				return err
-			}
-			defer f.Close()
-			reader = f
-		}
-
-		// Parse the Modelfile
-		modelfile, err := parser.ParseFile(reader)
-		if err != nil {
-			return fmt.Errorf("failed to parse Modelfile: %w", err)
-		}
-
-		// Extract FROM path and configuration
-		var modelDir string
-		mfConfig := &xcreateclient.ModelfileConfig{}
-
-		for _, cmd := range modelfile.Commands {
-			switch cmd.Name {
-			case "model":
-				modelDir = cmd.Args
-			case "template":
-				mfConfig.Template = cmd.Args
-			case "system":
-				mfConfig.System = cmd.Args
-			case "license":
-				mfConfig.License = cmd.Args
-			}
-		}
-
-		if modelDir == "" {
-			modelDir = "."
-		}
-
-		// Resolve relative paths based on Modelfile location
-		if !filepath.IsAbs(modelDir) && filename != "" {
-			modelDir = filepath.Join(filepath.Dir(filename), modelDir)
-		}
-
-		quantize, _ := cmd.Flags().GetString("quantize")
-		return xcreateclient.CreateModel(xcreateclient.CreateOptions{
-			ModelName: modelName,
-			ModelDir:  modelDir,
-			Quantize:  quantize,
-			Modelfile: mfConfig,
-		}, p)
-	}
-
 	var reader io.Reader
 
 	filename, err := getModelfileName(cmd)
 	if os.IsNotExist(err) {
 		if filename == "" {
 			// No Modelfile found - check if current directory is an image gen model
-			if create.IsTensorModelDir(".") {
+			if imagegen.IsTensorModelDir(".") {
 				quantize, _ := cmd.Flags().GetString("quantize")
-				return xcreateclient.CreateModel(xcreateclient.CreateOptions{
-					ModelName: modelName,
-					ModelDir:  ".",
-					Quantize:  quantize,
-				}, p)
+				return imagegenclient.CreateModel(args[0], ".", quantize, p)
 			}
 			reader = strings.NewReader("FROM .\n")
 		} else {
@@ -207,7 +134,7 @@ func CreateHandler(cmd *cobra.Command, args []string) error {
 	}
 	spinner.Stop()
 
-	req.Model = modelName
+	req.Model = args[0]
 	quantize, _ := cmd.Flags().GetString("quantize")
 	if quantize != "" {
 		req.Quantize = quantize
@@ -600,7 +527,7 @@ func RunHandler(cmd *cobra.Command, args []string) error {
 	}
 
 	// Check if this is an image generation model
-	if slices.Contains(info.Capabilities, model.CapabilityImage) {
+	if slices.Contains(info.Capabilities, model.CapabilityImageGeneration) {
 		if opts.Prompt == "" && !interactive {
 			return errors.New("image generation models require a prompt. Usage: ollama run " + name + " \"your prompt here\"")
 		}
@@ -1815,22 +1742,15 @@ func NewCLI() *cobra.Command {
 	rootCmd.Flags().BoolP("version", "v", false, "Show version information")
 
 	createCmd := &cobra.Command{
-		Use:   "create MODEL",
-		Short: "Create a model",
-		Args:  cobra.ExactArgs(1),
-		PreRunE: func(cmd *cobra.Command, args []string) error {
-			// Skip server check for experimental mode (writes directly to disk)
-			if experimental, _ := cmd.Flags().GetBool("experimental"); experimental {
-				return nil
-			}
-			return checkServerHeartbeat(cmd, args)
-		},
-		RunE: CreateHandler,
+		Use:     "create MODEL",
+		Short:   "Create a model",
+		Args:    cobra.ExactArgs(1),
+		PreRunE: checkServerHeartbeat,
+		RunE:    CreateHandler,
 	}
 
 	createCmd.Flags().StringP("file", "f", "", "Name of the Modelfile (default \"Modelfile\")")
 	createCmd.Flags().StringP("quantize", "q", "", "Quantize model to this level (e.g. q4_K_M)")
-	createCmd.Flags().Bool("experimental", false, "Enable experimental safetensors model creation")
 
 	showCmd := &cobra.Command{
 		Use:     "show MODEL",
@@ -1958,7 +1878,7 @@ func NewCLI() *cobra.Command {
 		Use:    "runner",
 		Hidden: true,
 		RunE: func(cmd *cobra.Command, args []string) error {
-			return runner.Execute(os.Args[2:])
+			return runner.Execute(os.Args[1:])
 		},
 		FParseErrWhitelist: cobra.FParseErrWhitelist{UnknownFlags: true},
 	}
@@ -1985,7 +1905,6 @@ func NewCLI() *cobra.Command {
 	} {
 		switch cmd {
 		case runCmd:
-			imagegen.AppendFlagsDocs(cmd)
 			appendEnvDocs(cmd, []envconfig.EnvVar{envVars["OLLAMA_HOST"], envVars["OLLAMA_NOHISTORY"]})
 		case serveCmd:
 			appendEnvDocs(cmd, []envconfig.EnvVar{

cmd/cmd_test.go

@@ -1555,7 +1555,7 @@ func TestShowInfoImageGen(t *testing.T) {
 			ParameterSize:     "10.3B",
 			QuantizationLevel: "FP8",
 		},
-		Capabilities: []model.Capability{model.CapabilityImage},
+		Capabilities: []model.Capability{model.CapabilityImageGeneration},
 		Requires:     "0.14.0",
 	}, false, &b)
 	if err != nil {

cmd/interactive.go

@@ -116,7 +116,7 @@ func generateInteractive(cmd *cobra.Command, opts runOptions) error {
 		Prompt:         ">>> ",
 		AltPrompt:      "... ",
 		Placeholder:    "Send a message (/? for help)",
-		AltPlaceholder: "Press Enter to send",
+		AltPlaceholder: `Use """ to end multi-line input`,
 	})
 	if err != nil {
 		return err

cmd/runner/main.go

@@ -0,0 +1,15 @@
+package main
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/ollama/ollama/runner"
+)
+
+func main() {
+	if err := runner.Execute(os.Args[1:]); err != nil {
+		fmt.Fprintf(os.Stderr, "error: %s\n", err)
+		os.Exit(1)
+	}
+}

docs/api.md

@@ -16,7 +16,6 @@
 - [Generate Embeddings](#generate-embeddings)
 - [List Running Models](#list-running-models)
 - [Version](#version)
-- [Experimental: Image Generation](#image-generation-experimental)
 
 ## Conventions
 
@@ -59,15 +58,6 @@ Advanced parameters (optional):
 - `keep_alive`: controls how long the model will stay loaded into memory following the request (default: `5m`)
 - `context` (deprecated): the context parameter returned from a previous request to `/generate`, this can be used to keep a short conversational memory
 
-Experimental image generation parameters (for image generation models only):
-
-> [!WARNING]
-> These parameters are experimental and may change in future versions.
-
-- `width`: width of the generated image in pixels
-- `height`: height of the generated image in pixels
-- `steps`: number of diffusion steps
-
 #### Structured outputs
 
 Structured outputs are supported by providing a JSON schema in the `format` parameter. The model will generate a response that matches the schema. See the [structured outputs](#request-structured-outputs) example below.
@@ -1877,55 +1867,3 @@ curl http://localhost:11434/api/version
   "version": "0.5.1"
 }
 ```
-
-## Experimental Features
-
-### Image Generation (Experimental)
-
-> [!WARNING]
-> Image generation is experimental and may change in future versions.
-
-Image generation is now supported through the standard `/api/generate` endpoint when using image generation models. The API automatically detects when an image generation model is being used.
-
-See the [Generate a completion](#generate-a-completion) section for the full API documentation. The experimental image generation parameters (`width`, `height`, `steps`) are documented there.
-
-#### Example
-
-##### Request
-
-```shell
-curl http://localhost:11434/api/generate -d '{
-  "model": "x/z-image-turbo",
-  "prompt": "a sunset over mountains",
-  "width": 1024,
-  "height": 768
-}'
-```
-
-##### Response (streaming)
-
-Progress updates during generation:
-
-```json
-{
-  "model": "x/z-image-turbo",
-  "created_at": "2024-01-15T10:30:00.000000Z",
-  "completed": 5,
-  "total": 20,
-  "done": false
-}
-```
-
-##### Final Response
-
-```json
-{
-  "model": "x/z-image-turbo",
-  "created_at": "2024-01-15T10:30:15.000000Z",
-  "image": "iVBORw0KGgoAAAANSUhEUg...",
-  "done": true,
-  "done_reason": "stop",
-  "total_duration": 15000000000,
-  "load_duration": 2000000000
-}
-```

docs/api/anthropic-compatibility.mdx

@@ -21,7 +21,6 @@ ollama pull glm-4.7:cloud
 To use Ollama with tools that expect the Anthropic API (like Claude Code), set these environment variables:
 
 ```shell
-export ANTHROPIC_AUTH_TOKEN=ollama  # required but ignored
 export ANTHROPIC_BASE_URL=http://localhost:11434
 export ANTHROPIC_API_KEY=ollama  # required but ignored
 ```
@@ -248,13 +247,12 @@ curl -X POST http://localhost:11434/v1/messages \
 [Claude Code](https://code.claude.com/docs/en/overview) can be configured to use Ollama as its backend:
 
 ```shell
-ANTHROPIC_AUTH_TOKEN=ollama ANTHROPIC_BASE_URL=http://localhost:11434 ANTHROPIC_API_KEY=ollama claude --model qwen3-coder
+ANTHROPIC_BASE_URL=http://localhost:11434 ANTHROPIC_API_KEY=ollama claude --model qwen3-coder
 ```
 
 Or set the environment variables in your shell profile:
 
 ```shell
-export ANTHROPIC_AUTH_TOKEN=ollama
 export ANTHROPIC_BASE_URL=http://localhost:11434
 export ANTHROPIC_API_KEY=ollama
 ```

docs/api/openai-compatibility.mdx

@@ -275,73 +275,6 @@ curl -X POST http://localhost:11434/v1/chat/completions \
 - [x] `dimensions`
 - [ ] `user`
 
-### `/v1/images/generations` (experimental)
-
-> Note: This endpoint is experimental and may change or be removed in future versions.
-
-Generate images using image generation models.
-
-<CodeGroup dropdown>
-
-```python images.py
-from openai import OpenAI
-
-client = OpenAI(
-    base_url='http://localhost:11434/v1/',
-    api_key='ollama',  # required but ignored
-)
-
-response = client.images.generate(
-    model='x/z-image-turbo',
-    prompt='A cute robot learning to paint',
-    size='1024x1024',
-    response_format='b64_json',
-)
-print(response.data[0].b64_json[:50] + '...')
-```
-
-```javascript images.js
-import OpenAI from "openai";
-
-const openai = new OpenAI({
-  baseURL: "http://localhost:11434/v1/",
-  apiKey: "ollama", // required but ignored
-});
-
-const response = await openai.images.generate({
-  model: "x/z-image-turbo",
-  prompt: "A cute robot learning to paint",
-  size: "1024x1024",
-  response_format: "b64_json",
-});
-
-console.log(response.data[0].b64_json.slice(0, 50) + "...");
-```
-
-```shell images.sh
-curl -X POST http://localhost:11434/v1/images/generations \
--H "Content-Type: application/json" \
--d '{
-  "model": "x/z-image-turbo",
-  "prompt": "A cute robot learning to paint",
-  "size": "1024x1024",
-  "response_format": "b64_json"
-}'
-```
-
-</CodeGroup>
-
-#### Supported request fields
-
-- [x] `model`
-- [x] `prompt`
-- [x] `size` (e.g. "1024x1024")
-- [x] `response_format` (only `b64_json` supported)
-- [ ] `n`
-- [ ] `quality`
-- [ ] `style`
-- [ ] `user`
-
 ### `/v1/responses`
 
 > Note: Added in Ollama v0.13.3

docs/capabilities/web-search.mdx

@@ -110,7 +110,7 @@ More Ollama [Python example](https://github.com/ollama/ollama-python/blob/main/e
 import { Ollama } from "ollama";
 
 const client = new Ollama();
-const results = await client.webSearch("what is ollama?");
+const results = await client.webSearch({ query: "what is ollama?" });
 console.log(JSON.stringify(results, null, 2));
 ```
 
@@ -213,7 +213,7 @@ models](https://ollama.com/models)\n\nAvailable for macOS, Windows, and Linux',
 import { Ollama } from "ollama";
 
 const client = new Ollama();
-const fetchResult = await client.webFetch("https://ollama.com");
+const fetchResult = await client.webFetch({ url: "https://ollama.com" });
 console.log(JSON.stringify(fetchResult, null, 2));
 ```
 

docs/docs.json

@@ -111,9 +111,7 @@
               "/integrations/zed",
               "/integrations/roo-code",
               "/integrations/n8n",
-              "/integrations/xcode",
-              "/integrations/onyx",
-              "/integrations/marimo"
+              "/integrations/xcode"
             ]
           },
           {

docs/faq.mdx

@@ -22,7 +22,7 @@ Please refer to the [GPU docs](./gpu).
 
 ## How can I specify the context window size?
 
-By default, Ollama uses a context window size of 4096 tokens.
+By default, Ollama uses a context window size of 2048 tokens.
 
 This can be overridden with the `OLLAMA_CONTEXT_LENGTH` environment variable. For example, to set the default context window to 8K, use:
 

docs/integrations/claude-code.mdx

@@ -2,12 +2,6 @@
 title: Claude Code
 ---
 
-Claude Code is Anthropic's agentic coding tool that can read, modify, and execute code in your working directory. 
-
-Open models can be used with Claude Code through Ollama's Anthropic-compatible API, enabling you to use models such as `qwen3-coder`, `gpt-oss:20b`, or other models.
-
-![Claude Code with Ollama](https://files.ollama.com/claude-code.png)
-
 ## Install
 
 Install [Claude Code](https://code.claude.com/docs/en/overview):
@@ -31,24 +25,22 @@ Claude Code connects to Ollama using the Anthropic-compatible API.
 1. Set the environment variables:
 
 ```shell
-export ANTHROPIC_AUTH_TOKEN=ollama
 export ANTHROPIC_BASE_URL=http://localhost:11434
+export ANTHROPIC_API_KEY=ollama
 ```
 
 2. Run Claude Code with an Ollama model:
 
 ```shell
-claude --model gpt-oss:20b
+claude --model qwen3-coder
 ```
 
 Or run with environment variables inline:
 
 ```shell
-ANTHROPIC_AUTH_TOKEN=ollama ANTHROPIC_BASE_URL=http://localhost:11434 claude --model gpt-oss:20b
+ANTHROPIC_BASE_URL=http://localhost:11434 ANTHROPIC_API_KEY=ollama claude --model qwen3-coder
 ```
 
-**Note:** Claude Code requires a large context window. We recommend at least 32K tokens. See the [context length documentation](/context-length) for how to adjust context length in Ollama.
-
 ## Connecting to ollama.com
 
 1. Create an [API key](https://ollama.com/settings/keys) on ollama.com
@@ -75,4 +67,3 @@ claude --model glm-4.7:cloud
 ### Local models
 - `qwen3-coder` - Excellent for coding tasks
 - `gpt-oss:20b` - Strong general-purpose model
-- `gpt-oss:120b` - Larger general-purpose model for more complex tasks

docs/integrations/marimo.mdx

@@ -1,73 +0,0 @@
----
-title: marimo
----
-
-## Install
-
-Install [marimo](https://marimo.io). You can use `pip` or `uv` for this. You 
-can also use `uv` to create a sandboxed environment for marimo by running:
-
-```
-uvx marimo edit --sandbox notebook.py
-```
-
-## Usage with Ollama
-
-1. In marimo, go to the user settings and go to the AI tab. From here
-you can find and configure Ollama as an AI provider. For local use you
-would typically point the base url to `http://localhost:11434/v1`.
-
-<div style={{ display: 'flex', justifyContent: 'center' }}>
-  <img 
-    src="/images/marimo-settings.png" 
-    alt="Ollama settings in marimo"
-    width="50%"
-  />
-</div>
-
-2. Once the AI provider is set up, you can turn on/off specific AI models you'd like to access. 
-
-<div style={{ display: 'flex', justifyContent: 'center' }}>
-  <img 
-    src="/images/marimo-models.png" 
-    alt="Selecting an Ollama model"
-    width="50%"
-  />
-</div>
-
-3. You can also add a model to the list of available models by scrolling to the bottom and using the UI there. 
-
-<div style={{ display: 'flex', justifyContent: 'center' }}>
-  <img 
-    src="/images/marimo-add-model.png" 
-    alt="Adding a new Ollama model"
-    width="50%"
-  />
-</div>
-
-4. Once configured, you can now use Ollama for AI chats in marimo.
-
-<div style={{ display: 'flex', justifyContent: 'center' }}>
-  <img 
-    src="/images/marimo-chat.png" 
-    alt="Configure code completion"
-    width="50%"
-  />
-</div>
-
-4. Alternatively, you can now use Ollama for **inline code completion** in marimo. This can be configured in the "AI Features" tab. 
-
-<div style={{ display: 'flex', justifyContent: 'center' }}>
-  <img 
-    src="/images/marimo-code-completion.png" 
-    alt="Configure code completion"
-    width="50%"
-  />
-</div>
-
-
-## Connecting to ollama.com
-
-1. Sign in to ollama cloud via `ollama signin` 
-2. In the ollama model settings add a model that ollama hosts, like `gpt-oss:120b`.
-3. You can now refer to this model in marimo!

docs/integrations/onyx.mdx

@@ -1,63 +0,0 @@
----
-title: Onyx
----
-
-## Overview
-[Onyx](http://onyx.app/) is a self-hostable Chat UI that integrates with all Ollama models. Features include:
-- Creating custom Agents
-- Web search
-- Deep Research
-- RAG over uploaded documents and connected apps
-- Connectors to applications like Google Drive, Email, Slack, etc.
-- MCP and OpenAPI Actions support
-- Image generation
-- User/Groups management, RBAC, SSO, etc.
-
-Onyx can be deployed for single users or large organizations.
-
-## Install Onyx
-
-Deploy Onyx with the [quickstart guide](https://docs.onyx.app/deployment/getting_started/quickstart).
-
-<Info>
-Resourcing/scaling docs [here](https://docs.onyx.app/deployment/getting_started/resourcing).
-</Info>
-
-## Usage with Ollama 
-
-1. Login to your Onyx deployment (create an account first).
-<div style={{ display: 'flex', justifyContent: 'center' }}>
-  <img 
-    src="/images/onyx-login.png" 
-    alt="Onyx Login Page"
-    width="75%"
-  />
-</div>
-2. In the set-up process select `Ollama` as the LLM provider.
-<div style={{ display: 'flex', justifyContent: 'center' }}>
-  <img 
-    src="/images/onyx-ollama-llm.png" 
-    alt="Onyx Set Up Form"
-    width="75%"
-  />
-</div>
-3. Provide your **Ollama API URL** and select your models.
-<Note>If you're running Onyx in Docker, to access your computer's local network use `http://host.docker.internal` instead of `http://127.0.0.1`.</Note>
-<div style={{ display: 'flex', justifyContent: 'center' }}>
-  <img 
-    src="/images/onyx-ollama-form.png" 
-    alt="Selecting Ollama Models"
-    width="75%"
-  />
-</div>
-
-You can also easily connect up Onyx Cloud with the `Ollama Cloud` tab of the setup.
-
-## Send your first query
-<div style={{ display: 'flex', justifyContent: 'center' }}>
-  <img 
-    src="/images/onyx-query.png" 
-    alt="Onyx Query Example"
-    width="75%"
-  />
-</div>

docs/linux.mdx

@@ -1,5 +1,5 @@
 ---
-title: Linux
+title: "Linux"
 ---
 
 ## Install
@@ -13,15 +13,14 @@ curl -fsSL https://ollama.com/install.sh | sh
 ## Manual install
 
 <Note>
-  If you are upgrading from a prior version, you should remove the old libraries
-  with `sudo rm -rf /usr/lib/ollama` first.
+  If you are upgrading from a prior version, you should remove the old libraries with `sudo rm -rf /usr/lib/ollama` first.
 </Note>
 
 Download and extract the package:
 
 ```shell
-curl -fsSL https://ollama.com/download/ollama-linux-amd64.tar.zst \
-    | sudo tar x -C /usr
+curl -fsSL https://ollama.com/download/ollama-linux-amd64.tgz \
+    | sudo tar zx -C /usr
 ```
 
 Start Ollama:
@@ -41,8 +40,8 @@ ollama -v
 If you have an AMD GPU, also download and extract the additional ROCm package:
 
 ```shell
-curl -fsSL https://ollama.com/download/ollama-linux-amd64-rocm.tar.zst \
-    | sudo tar x -C /usr
+curl -fsSL https://ollama.com/download/ollama-linux-amd64-rocm.tgz \
+    | sudo tar zx -C /usr
 ```
 
 ### ARM64 install
@@ -50,8 +49,8 @@ curl -fsSL https://ollama.com/download/ollama-linux-amd64-rocm.tar.zst \
 Download and extract the ARM64-specific package:
 
 ```shell
-curl -fsSL https://ollama.com/download/ollama-linux-arm64.tar.zst \
-    | sudo tar x -C /usr
+curl -fsSL https://ollama.com/download/ollama-linux-arm64.tgz \
+    | sudo tar zx -C /usr
 ```
 
 ### Adding Ollama as a startup service (recommended)
@@ -113,11 +112,7 @@ sudo systemctl status ollama
 ```
 
 <Note>
-  While AMD has contributed the `amdgpu` driver upstream to the official linux
-  kernel source, the version is older and may not support all ROCm features. We
-  recommend you install the latest driver from
-  https://www.amd.com/en/support/linux-drivers for best support of your Radeon
-  GPU.
+  While AMD has contributed the `amdgpu` driver upstream to the official linux kernel source, the version is older and may not support all ROCm features. We recommend you install the latest driver from https://www.amd.com/en/support/linux-drivers for best support of your Radeon GPU.
 </Note>
 
 ## Customizing
@@ -146,8 +141,8 @@ curl -fsSL https://ollama.com/install.sh | sh
 Or by re-downloading Ollama:
 
 ```shell
-curl -fsSL https://ollama.com/download/ollama-linux-amd64.tar.zst \
-    | sudo tar x -C /usr
+curl -fsSL https://ollama.com/download/ollama-linux-amd64.tgz \
+    | sudo tar zx -C /usr
 ```
 
 ## Installing specific versions
@@ -196,4 +191,4 @@ Remove the downloaded models and Ollama service user and group:
 sudo userdel ollama
 sudo groupdel ollama
 sudo rm -r /usr/share/ollama
-```
+```

integration/tools_test.go

@@ -131,7 +131,7 @@ func TestAPIToolCalling(t *testing.T) {
 					t.Errorf("unexpected tool called: got %q want %q", lastToolCall.Function.Name, "get_weather")
 				}
 
-				if _, ok := lastToolCall.Function.Arguments.Get("location"); !ok {
+				if _, ok := lastToolCall.Function.Arguments["location"]; !ok {
 					t.Errorf("expected tool arguments to include 'location', got: %s", lastToolCall.Function.Arguments.String())
 				}
 			case <-ctx.Done():

llm/server.go

@@ -1464,12 +1464,6 @@ type CompletionRequest struct {
 
 	// TopLogprobs specifies the number of most likely alternative tokens to return (0-20)
 	TopLogprobs int
-
-	// Image generation fields
-	Width  int32 `json:"width,omitempty"`
-	Height int32 `json:"height,omitempty"`
-	Steps  int32 `json:"steps,omitempty"`
-	Seed   int64 `json:"seed,omitempty"`
 }
 
 // DoneReason represents the reason why a completion response is done
@@ -1518,15 +1512,6 @@ type CompletionResponse struct {
 
 	// Logprobs contains log probability information if requested
 	Logprobs []Logprob `json:"logprobs,omitempty"`
-
-	// Image contains base64-encoded image data for image generation
-	Image string `json:"image,omitempty"`
-
-	// Step is the current step in image generation
-	Step int `json:"step,omitempty"`
-
-	// TotalSteps is the total number of steps for image generation
-	TotalSteps int `json:"total_steps,omitempty"`
 }
 
 func (s *llmServer) Completion(ctx context.Context, req CompletionRequest, fn func(CompletionResponse)) error {

middleware/openai.go

@@ -8,7 +8,6 @@ import (
 	"math/rand"
 	"net/http"
 	"strings"
-	"time"
 
 	"github.com/gin-gonic/gin"
 
@@ -442,7 +441,6 @@ type ResponsesWriter struct {
 	stream     bool
 	responseID string
 	itemID     string
-	request    openai.ResponsesRequest
 }
 
 func (w *ResponsesWriter) writeEvent(eventType string, data any) error {
@@ -480,9 +478,7 @@ func (w *ResponsesWriter) writeResponse(data []byte) (int, error) {
 
 	// Non-streaming response
 	w.ResponseWriter.Header().Set("Content-Type", "application/json")
-	response := openai.ToResponse(w.model, w.responseID, w.itemID, chatResponse, w.request)
-	completedAt := time.Now().Unix()
-	response.CompletedAt = &completedAt
+	response := openai.ToResponse(w.model, w.responseID, w.itemID, chatResponse)
 	return len(data), json.NewEncoder(w.ResponseWriter).Encode(response)
 }
 
@@ -527,12 +523,11 @@ func ResponsesMiddleware() gin.HandlerFunc {
 
 		w := &ResponsesWriter{
 			BaseWriter: BaseWriter{ResponseWriter: c.Writer},
-			converter:  openai.NewResponsesStreamConverter(responseID, itemID, req.Model, req),
+			converter:  openai.NewResponsesStreamConverter(responseID, itemID, req.Model),
 			model:      req.Model,
 			stream:     streamRequested,
 			responseID: responseID,
 			itemID:     itemID,
-			request:    req,
 		}
 
 		// Set headers based on streaming mode
@@ -546,66 +541,3 @@ func ResponsesMiddleware() gin.HandlerFunc {
 		c.Next()
 	}
 }
-
-type ImageWriter struct {
-	BaseWriter
-}
-
-func (w *ImageWriter) writeResponse(data []byte) (int, error) {
-	var generateResponse api.GenerateResponse
-	if err := json.Unmarshal(data, &generateResponse); err != nil {
-		return 0, err
-	}
-
-	// Only write response when done with image
-	if generateResponse.Done && generateResponse.Image != "" {
-		w.ResponseWriter.Header().Set("Content-Type", "application/json")
-		return len(data), json.NewEncoder(w.ResponseWriter).Encode(openai.ToImageGenerationResponse(generateResponse))
-	}
-
-	return len(data), nil
-}
-
-func (w *ImageWriter) Write(data []byte) (int, error) {
-	code := w.ResponseWriter.Status()
-	if code != http.StatusOK {
-		return w.writeError(data)
-	}
-
-	return w.writeResponse(data)
-}
-
-func ImageGenerationsMiddleware() gin.HandlerFunc {
-	return func(c *gin.Context) {
-		var req openai.ImageGenerationRequest
-		if err := c.ShouldBindJSON(&req); err != nil {
-			c.AbortWithStatusJSON(http.StatusBadRequest, openai.NewError(http.StatusBadRequest, err.Error()))
-			return
-		}
-
-		if req.Prompt == "" {
-			c.AbortWithStatusJSON(http.StatusBadRequest, openai.NewError(http.StatusBadRequest, "prompt is required"))
-			return
-		}
-
-		if req.Model == "" {
-			c.AbortWithStatusJSON(http.StatusBadRequest, openai.NewError(http.StatusBadRequest, "model is required"))
-			return
-		}
-
-		var b bytes.Buffer
-		if err := json.NewEncoder(&b).Encode(openai.FromImageGenerationRequest(req)); err != nil {
-			c.AbortWithStatusJSON(http.StatusInternalServerError, openai.NewError(http.StatusInternalServerError, err.Error()))
-			return
-		}
-
-		c.Request.Body = io.NopCloser(&b)
-
-		w := &ImageWriter{
-			BaseWriter: BaseWriter{ResponseWriter: c.Writer},
-		}
-
-		c.Writer = w
-		c.Next()
-	}
-}

middleware/openai_test.go

@@ -961,154 +961,3 @@ func TestRetrieveMiddleware(t *testing.T) {
 		}
 	}
 }
-
-func TestImageGenerationsMiddleware(t *testing.T) {
-	type testCase struct {
-		name string
-		body string
-		req  api.GenerateRequest
-		err  openai.ErrorResponse
-	}
-
-	var capturedRequest *api.GenerateRequest
-
-	testCases := []testCase{
-		{
-			name: "image generation basic",
-			body: `{
-				"model": "test-model",
-				"prompt": "a beautiful sunset"
-			}`,
-			req: api.GenerateRequest{
-				Model:  "test-model",
-				Prompt: "a beautiful sunset",
-			},
-		},
-		{
-			name: "image generation with size",
-			body: `{
-				"model": "test-model",
-				"prompt": "a beautiful sunset",
-				"size": "512x768"
-			}`,
-			req: api.GenerateRequest{
-				Model:  "test-model",
-				Prompt: "a beautiful sunset",
-				Width:  512,
-				Height: 768,
-			},
-		},
-		{
-			name: "image generation missing prompt",
-			body: `{
-				"model": "test-model"
-			}`,
-			err: openai.ErrorResponse{
-				Error: openai.Error{
-					Message: "prompt is required",
-					Type:    "invalid_request_error",
-				},
-			},
-		},
-		{
-			name: "image generation missing model",
-			body: `{
-				"prompt": "a beautiful sunset"
-			}`,
-			err: openai.ErrorResponse{
-				Error: openai.Error{
-					Message: "model is required",
-					Type:    "invalid_request_error",
-				},
-			},
-		},
-	}
-
-	endpoint := func(c *gin.Context) {
-		c.Status(http.StatusOK)
-	}
-
-	gin.SetMode(gin.TestMode)
-	router := gin.New()
-	router.Use(ImageGenerationsMiddleware(), captureRequestMiddleware(&capturedRequest))
-	router.Handle(http.MethodPost, "/api/generate", endpoint)
-
-	for _, tc := range testCases {
-		t.Run(tc.name, func(t *testing.T) {
-			req, _ := http.NewRequest(http.MethodPost, "/api/generate", strings.NewReader(tc.body))
-			req.Header.Set("Content-Type", "application/json")
-
-			defer func() { capturedRequest = nil }()
-
-			resp := httptest.NewRecorder()
-			router.ServeHTTP(resp, req)
-
-			if tc.err.Error.Message != "" {
-				var errResp openai.ErrorResponse
-				if err := json.Unmarshal(resp.Body.Bytes(), &errResp); err != nil {
-					t.Fatal(err)
-				}
-				if diff := cmp.Diff(tc.err, errResp); diff != "" {
-					t.Fatalf("errors did not match:\n%s", diff)
-				}
-				return
-			}
-
-			if resp.Code != http.StatusOK {
-				t.Fatalf("expected status 200, got %d: %s", resp.Code, resp.Body.String())
-			}
-
-			if diff := cmp.Diff(&tc.req, capturedRequest); diff != "" {
-				t.Fatalf("requests did not match:\n%s", diff)
-			}
-		})
-	}
-}
-
-func TestImageWriterResponse(t *testing.T) {
-	gin.SetMode(gin.TestMode)
-
-	// Test that ImageWriter transforms GenerateResponse to OpenAI format
-	endpoint := func(c *gin.Context) {
-		resp := api.GenerateResponse{
-			Model:     "test-model",
-			CreatedAt: time.Unix(1234567890, 0).UTC(),
-			Done:      true,
-			Image:     "dGVzdC1pbWFnZS1kYXRh", // base64 of "test-image-data"
-		}
-		data, _ := json.Marshal(resp)
-		c.Writer.Write(append(data, '\n'))
-	}
-
-	router := gin.New()
-	router.Use(ImageGenerationsMiddleware())
-	router.Handle(http.MethodPost, "/api/generate", endpoint)
-
-	body := `{"model": "test-model", "prompt": "test"}`
-	req, _ := http.NewRequest(http.MethodPost, "/api/generate", strings.NewReader(body))
-	req.Header.Set("Content-Type", "application/json")
-
-	resp := httptest.NewRecorder()
-	router.ServeHTTP(resp, req)
-
-	if resp.Code != http.StatusOK {
-		t.Fatalf("expected status 200, got %d: %s", resp.Code, resp.Body.String())
-	}
-
-	var imageResp openai.ImageGenerationResponse
-	if err := json.Unmarshal(resp.Body.Bytes(), &imageResp); err != nil {
-		t.Fatalf("failed to unmarshal response: %v", err)
-	}
-
-	if imageResp.Created != 1234567890 {
-		t.Errorf("expected created 1234567890, got %d", imageResp.Created)
-	}
-
-	if len(imageResp.Data) != 1 {
-		t.Fatalf("expected 1 image, got %d", len(imageResp.Data))
-	}
-
-	if imageResp.Data[0].B64JSON != "dGVzdC1pbWFnZS1kYXRh" {
-		t.Errorf("expected image data 'dGVzdC1pbWFnZS1kYXRh', got %s", imageResp.Data[0].B64JSON)
-	}
-}

model/parsers/nemotron3nano.go

@@ -1,6 +1,7 @@
 package parsers
 
 import (
+	"regexp"
 	"strings"
 	"unicode"
 
@@ -13,114 +14,243 @@ const (
 	Nemotron3NanoCollectingThinking Nemotron3NanoParserState = iota
 	Nemotron3NanoSkipWhitespaceAfterThinking
 	Nemotron3NanoCollectingContent
+	Nemotron3NanoCollectingToolCalls
 )
 
 const (
-	nemotronThinkClose   = "</think>"
-	nemotronToolCallOpen = "<tool_call>"
+	nemotronThinkClose    = "</think>"
+	nemotronToolCallOpen  = "<tool_call>"
+	nemotronToolCallClose = "</tool_call>"
 )
 
 type Nemotron3NanoParser struct {
-	state      Nemotron3NanoParserState
-	buffer     strings.Builder
-	toolParser *Qwen3CoderParser
+	state  Nemotron3NanoParserState
+	buffer strings.Builder
+	tools  []api.Tool
 }
 
 func (p *Nemotron3NanoParser) HasToolSupport() bool     { return true }
 func (p *Nemotron3NanoParser) HasThinkingSupport() bool { return true }
 
 func (p *Nemotron3NanoParser) Init(tools []api.Tool, lastMessage *api.Message, thinkValue *api.ThinkValue) []api.Tool {
-	p.toolParser = &Qwen3CoderParser{}
-	p.toolParser.Init(tools, nil, nil)
+	p.tools = tools
 
+	// thinking is enabled if user requests it
 	thinkingEnabled := thinkValue != nil && thinkValue.Bool()
+
 	prefill := lastMessage != nil && lastMessage.Role == "assistant"
 
-	if !thinkingEnabled || (prefill && lastMessage.Content != "") {
+	if !thinkingEnabled {
 		p.state = Nemotron3NanoCollectingContent
-	} else {
-		p.state = Nemotron3NanoCollectingThinking
+		return tools
 	}
 
+	if prefill && lastMessage.Content != "" {
+		p.state = Nemotron3NanoCollectingContent
+		return tools
+	}
+
+	p.state = Nemotron3NanoCollectingThinking
 	return tools
 }
 
+type nemotronEvent interface {
+	isNemotronEvent()
+}
+
+type nemotronEventThinkingContent struct {
+	content string
+}
+
+type nemotronEventContent struct {
+	content string
+}
+
+type nemotronEventToolCall struct {
+	toolCall api.ToolCall
+}
+
+func (nemotronEventThinkingContent) isNemotronEvent() {}
+func (nemotronEventContent) isNemotronEvent()         {}
+func (nemotronEventToolCall) isNemotronEvent()        {}
+
 func (p *Nemotron3NanoParser) Add(s string, done bool) (content string, thinking string, calls []api.ToolCall, err error) {
-	if p.state == Nemotron3NanoCollectingContent {
-		return p.toolParser.Add(s, done)
+	p.buffer.WriteString(s)
+	events := p.parseEvents()
+
+	var toolCalls []api.ToolCall
+	var contentSb strings.Builder
+	var thinkingSb strings.Builder
+	for _, event := range events {
+		switch event := event.(type) {
+		case nemotronEventToolCall:
+			toolCalls = append(toolCalls, event.toolCall)
+		case nemotronEventThinkingContent:
+			thinkingSb.WriteString(event.content)
+		case nemotronEventContent:
+			contentSb.WriteString(event.content)
+		}
 	}
 
-	if p.state == Nemotron3NanoSkipWhitespaceAfterThinking {
-		s = strings.TrimLeftFunc(s, unicode.IsSpace)
-		if s == "" {
-			return "", "", nil, nil
+	return contentSb.String(), thinkingSb.String(), toolCalls, nil
+}
+
+func (p *Nemotron3NanoParser) parseEvents() []nemotronEvent {
+	var all []nemotronEvent
+
+	keepLooping := true
+	for keepLooping {
+		var events []nemotronEvent
+		events, keepLooping = p.eat()
+		if len(events) > 0 {
+			all = append(all, events...)
+		}
+	}
+
+	return all
+}
+
+// emitWithPartialCheck extracts unambiguous content before a potential partial tag
+func (p *Nemotron3NanoParser) emitWithPartialCheck(bufStr, tag string) (unambiguous, ambiguous string) {
+	if overlapLen := overlap(bufStr, tag); overlapLen > 0 {
+		beforePartialTag := bufStr[:len(bufStr)-overlapLen]
+		trailingLen := trailingWhitespaceLen(beforePartialTag)
+		return bufStr[:len(beforePartialTag)-trailingLen], bufStr[len(beforePartialTag)-trailingLen:]
+	}
+	wsLen := trailingWhitespaceLen(bufStr)
+	return bufStr[:len(bufStr)-wsLen], bufStr[len(bufStr)-wsLen:]
+}
+
+func (p *Nemotron3NanoParser) eat() ([]nemotronEvent, bool) {
+	bufStr := p.buffer.String()
+	if bufStr == "" {
+		return nil, false
+	}
+
+	switch p.state {
+	case Nemotron3NanoCollectingThinking:
+		if strings.Contains(bufStr, nemotronThinkClose) {
+			split := strings.SplitN(bufStr, nemotronThinkClose, 2)
+			thinking := strings.TrimRightFunc(split[0], unicode.IsSpace)
+			p.buffer.Reset()
+			remainder := strings.TrimLeftFunc(split[1], unicode.IsSpace)
+			p.buffer.WriteString(remainder)
+			// Transition to whitespace-skipping state if buffer is empty,
+			// otherwise go directly to content collection
+			if remainder == "" {
+				p.state = Nemotron3NanoSkipWhitespaceAfterThinking
+			} else {
+				p.state = Nemotron3NanoCollectingContent
+			}
+			if thinking != "" {
+				return []nemotronEvent{nemotronEventThinkingContent{content: thinking}}, true
+			}
+			return nil, true
+		}
+		unambig, ambig := p.emitWithPartialCheck(bufStr, nemotronThinkClose)
+		p.buffer.Reset()
+		p.buffer.WriteString(ambig)
+		if unambig != "" {
+			return []nemotronEvent{nemotronEventThinkingContent{content: unambig}}, false
+		}
+		return nil, false
+
+	// We only want to skip whitespace between thinking and content
+	case Nemotron3NanoSkipWhitespaceAfterThinking:
+		bufStr = strings.TrimLeftFunc(bufStr, unicode.IsSpace)
+		p.buffer.Reset()
+		p.buffer.WriteString(bufStr)
+		if bufStr == "" {
+			return nil, false
 		}
 		p.state = Nemotron3NanoCollectingContent
-		return p.toolParser.Add(s, done)
-	}
+		return nil, true
 
-	// Nemotron3NanoCollectingThinking - buffer and look for end markers
-	p.buffer.WriteString(s)
-	bufStr := p.buffer.String()
-
-	// Look for end of thinking: </think> or <tool_call> (model may skip </think>)
-	thinkIdx := strings.Index(bufStr, nemotronThinkClose)
-	toolIdx := strings.Index(bufStr, nemotronToolCallOpen)
-
-	var endIdx int = -1
-	var remainder string
-
-	if thinkIdx != -1 && (toolIdx == -1 || thinkIdx < toolIdx) {
-		endIdx = thinkIdx
-		remainder = strings.TrimLeftFunc(bufStr[thinkIdx+len(nemotronThinkClose):], unicode.IsSpace)
-	} else if toolIdx != -1 {
-		endIdx = toolIdx
-		remainder = bufStr[toolIdx:] // Include <tool_call> tag
-	}
-
-	if endIdx != -1 {
-		thinking = strings.TrimRightFunc(bufStr[:endIdx], unicode.IsSpace)
-		p.buffer.Reset()
-
-		if remainder == "" {
-			p.state = Nemotron3NanoSkipWhitespaceAfterThinking
-		} else {
-			p.state = Nemotron3NanoCollectingContent
-			content, _, calls, err = p.toolParser.Add(remainder, done)
+	case Nemotron3NanoCollectingContent:
+		if strings.Contains(bufStr, nemotronToolCallOpen) {
+			split := strings.SplitN(bufStr, nemotronToolCallOpen, 2)
+			content := strings.TrimRightFunc(split[0], unicode.IsSpace)
+			p.buffer.Reset()
+			p.buffer.WriteString(split[1])
+			p.state = Nemotron3NanoCollectingToolCalls
+			if content != "" {
+				return []nemotronEvent{nemotronEventContent{content: content}}, true
+			}
+			return nil, true
 		}
-		return content, thinking, calls, err
+		unambig, ambig := p.emitWithPartialCheck(bufStr, nemotronToolCallOpen)
+		p.buffer.Reset()
+		p.buffer.WriteString(ambig)
+		if unambig != "" {
+			return []nemotronEvent{nemotronEventContent{content: unambig}}, false
+		}
+		return nil, false
+
+	case Nemotron3NanoCollectingToolCalls:
+		if strings.Contains(bufStr, nemotronToolCallClose) {
+			split := strings.SplitN(bufStr, nemotronToolCallClose, 2)
+			remaining := strings.TrimLeftFunc(split[1], unicode.IsSpace)
+			p.buffer.Reset()
+			p.buffer.WriteString(remaining)
+
+			var events []nemotronEvent
+			if tc, err := p.parseToolCall(split[0]); err == nil {
+				events = append(events, nemotronEventToolCall{toolCall: tc})
+			}
+
+			if !strings.Contains(remaining, nemotronToolCallOpen) {
+				p.state = Nemotron3NanoCollectingContent
+			}
+			return events, true
+		}
+		return nil, false
 	}
 
-	// No end marker - emit unambiguous thinking
-	thinking = p.emitThinking(bufStr)
-	return "", thinking, nil, nil
+	return nil, false
 }
 
-// emitThinking returns unambiguous thinking content, keeping potential partial tags in buffer
-func (p *Nemotron3NanoParser) emitThinking(bufStr string) string {
-	// Check for partial </think> or <tool_call> at end
-	thinkOverlap := overlap(bufStr, nemotronThinkClose)
-	toolOverlap := overlap(bufStr, nemotronToolCallOpen)
-	maxOverlap := max(thinkOverlap, toolOverlap)
+var (
+	nemotronFunctionRegex  = regexp.MustCompile(`<function=([^>]+)>`)
+	nemotronParameterRegex = regexp.MustCompile(`<parameter=([^>]+)>\n?([\s\S]*?)\n?</parameter>`)
+)
 
-	if maxOverlap > 0 {
-		unambiguous := bufStr[:len(bufStr)-maxOverlap]
-		unambiguous = strings.TrimRightFunc(unambiguous, unicode.IsSpace)
-		p.buffer.Reset()
-		p.buffer.WriteString(bufStr[len(bufStr)-maxOverlap:])
-		return unambiguous
+func (p *Nemotron3NanoParser) parseToolCall(content string) (api.ToolCall, error) {
+	toolCall := api.ToolCall{}
+
+	// Extract function name
+	fnMatch := nemotronFunctionRegex.FindStringSubmatch(content)
+	if len(fnMatch) < 2 {
+		return toolCall, nil
+	}
+	toolCall.Function.Name = fnMatch[1]
+
+	// Extract parameters
+	toolCall.Function.Arguments = api.NewToolCallFunctionArguments()
+	paramMatches := nemotronParameterRegex.FindAllStringSubmatch(content, -1)
+	for _, match := range paramMatches {
+		if len(match) >= 3 {
+			paramName := match[1]
+			paramValue := strings.TrimSpace(match[2])
+
+			// Try to parse as typed value based on tool definition
+			toolCall.Function.Arguments.Set(paramName, p.parseParamValue(paramName, paramValue))
+		}
 	}
 
-	// No partial tags - emit all but trailing whitespace
-	wsLen := trailingWhitespaceLen(bufStr)
-	if wsLen > 0 {
-		unambiguous := bufStr[:len(bufStr)-wsLen]
-		p.buffer.Reset()
-		p.buffer.WriteString(bufStr[len(bufStr)-wsLen:])
-		return unambiguous
-	}
-
-	// Nothing to hold back
-	p.buffer.Reset()
-	return bufStr
+	return toolCall, nil
+}
+
+func (p *Nemotron3NanoParser) parseParamValue(paramName string, raw string) any {
+	// Find the matching tool to get parameter type
+	var paramType api.PropertyType
+	for _, tool := range p.tools {
+		if tool.Function.Parameters.Properties != nil {
+			if prop, ok := tool.Function.Parameters.Properties.Get(paramName); ok {
+				paramType = prop.Type
+				break
+			}
+		}
+	}
+
+	return parseValue(raw, paramType)
 }

model/parsers/nemotron3nano_test.go

@@ -8,8 +8,6 @@ import (
 	"github.com/ollama/ollama/api"
 )
 
-// TestNemotron3NanoParser tests Nemotron-specific behavior (thinking support).
-// Tool call parsing is tested in qwen3coder_test.go since Nemotron delegates to Qwen3CoderParser.
 func TestNemotron3NanoParser(t *testing.T) {
 	tests := []struct {
 		name             string
@@ -19,6 +17,18 @@ func TestNemotron3NanoParser(t *testing.T) {
 		expectedThinking string
 		expectedCalls    []api.ToolCall
 	}{
+		{
+			name:            "simple content - no thinking",
+			input:           "Hello, how can I help you?",
+			thinkValue:      nil,
+			expectedContent: "Hello, how can I help you?",
+		},
+		{
+			name:            "simple content - thinking disabled",
+			input:           "Hello, how can I help you?",
+			thinkValue:      &api.ThinkValue{Value: false},
+			expectedContent: "Hello, how can I help you?",
+		},
 		{
 			name:             "thinking then content",
 			input:            "Let me think about this...</think>\nHere is my answer.",
@@ -33,6 +43,69 @@ func TestNemotron3NanoParser(t *testing.T) {
 			expectedThinking: "Step 1: Analyze\nStep 2: Process\nStep 3: Conclude",
 			expectedContent:  "The answer is 42.",
 		},
+		{
+			name:       "simple tool call",
+			input:      "<tool_call>\n<function=get_weather>\n<parameter=city>\nParis\n</parameter>\n</function>\n</tool_call>",
+			thinkValue: nil,
+			expectedCalls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Name:      "get_weather",
+						Arguments: testArgs(map[string]any{"city": "Paris"}),
+					},
+				},
+			},
+		},
+		{
+			name:            "content then tool call",
+			input:           "Let me check the weather.\n<tool_call>\n<function=get_weather>\n<parameter=city>\nNYC\n</parameter>\n</function>\n</tool_call>",
+			thinkValue:      nil,
+			expectedContent: "Let me check the weather.",
+			expectedCalls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Name:      "get_weather",
+						Arguments: testArgs(map[string]any{"city": "NYC"}),
+					},
+				},
+			},
+		},
+		{
+			name:       "tool call with multiple parameters",
+			input:      "<tool_call>\n<function=book_flight>\n<parameter=from>\nSFO\n</parameter>\n<parameter=to>\nNYC\n</parameter>\n</function>\n</tool_call>",
+			thinkValue: nil,
+			expectedCalls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Name: "book_flight",
+						Arguments: testArgs(map[string]any{
+							"from": "SFO",
+							"to":   "NYC",
+						}),
+					},
+				},
+			},
+		},
+		{
+			name: "multiple tool calls",
+			input: "<tool_call>\n<function=get_weather>\n<parameter=city>\nSan Francisco\n</parameter>\n</function>\n</tool_call>\n" +
+				"<tool_call>\n<function=get_weather>\n<parameter=city>\nNew York\n</parameter>\n</function>\n</tool_call>",
+			thinkValue: nil,
+			expectedCalls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Name:      "get_weather",
+						Arguments: testArgs(map[string]any{"city": "San Francisco"}),
+					},
+				},
+				{
+					Function: api.ToolCallFunction{
+						Name:      "get_weather",
+						Arguments: testArgs(map[string]any{"city": "New York"}),
+					},
+				},
+			},
+		},
 		{
 			name:             "thinking then tool call",
 			input:            "I should check the weather...</think>\n<tool_call>\n<function=get_weather>\n<parameter=city>\nParis\n</parameter>\n</function>\n</tool_call>",
@@ -62,6 +135,19 @@ func TestNemotron3NanoParser(t *testing.T) {
 				},
 			},
 		},
+		{
+			name:       "tool call with multiline parameter value",
+			input:      "<tool_call>\n<function=create_note>\n<parameter=content>\nLine 1\nLine 2\nLine 3\n</parameter>\n</function>\n</tool_call>",
+			thinkValue: nil,
+			expectedCalls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Name:      "create_note",
+						Arguments: testArgs(map[string]any{"content": "Line 1\nLine 2\nLine 3"}),
+					},
+				},
+			},
+		},
 		{
 			name:             "empty thinking block - immediate close",
 			input:            "</think>\nHere is my answer.",
@@ -75,6 +161,18 @@ func TestNemotron3NanoParser(t *testing.T) {
 			thinkValue:      &api.ThinkValue{Value: false},
 			expectedContent: "</think>\nSome content after spurious tag.",
 		},
+		{
+			name:          "tool call with no function name - returns empty tool call",
+			input:         "<tool_call>\n<function=>\n</function>\n</tool_call>",
+			thinkValue:    nil,
+			expectedCalls: []api.ToolCall{{Function: api.ToolCallFunction{Name: "", Arguments: api.NewToolCallFunctionArguments()}}},
+		},
+		{
+			name:            "content with newlines preserved",
+			input:           "Line 1\n\nLine 2\n\n\nLine 3",
+			thinkValue:      nil,
+			expectedContent: "Line 1\n\nLine 2\n\n\nLine 3",
+		},
 		{
 			name:             "thinking with only whitespace after close tag",
 			input:            "My thoughts...</think>   \n\t\n   Content here.",
@@ -82,6 +180,25 @@ func TestNemotron3NanoParser(t *testing.T) {
 			expectedThinking: "My thoughts...",
 			expectedContent:  "Content here.",
 		},
+		{
+			name:            "unicode content",
+			input:           "Hello 世界! 🌍 Ñoño",
+			thinkValue:      nil,
+			expectedContent: "Hello 世界! 🌍 Ñoño",
+		},
+		{
+			name:       "tool call with numeric parameter",
+			input:      "<tool_call>\n<function=set_temp>\n<parameter=value>\n42\n</parameter>\n</function>\n</tool_call>",
+			thinkValue: nil,
+			expectedCalls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Name:      "set_temp",
+						Arguments: testArgs(map[string]any{"value": "42"}),
+					},
+				},
+			},
+		},
 	}
 
 	for _, tt := range tests {
@@ -116,8 +233,6 @@ func TestNemotron3NanoParser(t *testing.T) {
 	}
 }
 
-// TestNemotron3NanoParser_Streaming tests streaming behavior for thinking support.
-// Tool call streaming is tested in qwen3coder_test.go.
 func TestNemotron3NanoParser_Streaming(t *testing.T) {
 	tests := []struct {
 		name             string
@@ -127,6 +242,18 @@ func TestNemotron3NanoParser_Streaming(t *testing.T) {
 		expectedThinking string
 		expectedCalls    []api.ToolCall
 	}{
+		{
+			name:            "streaming content character by character",
+			chunks:          []string{"H", "e", "l", "l", "o", ",", " ", "w", "o", "r", "l", "d", "!"},
+			thinkValue:      nil,
+			expectedContent: "Hello, world!",
+		},
+		{
+			name:            "streaming content small tokens",
+			chunks:          []string{"Hel", "lo", ", ", "how ", "can", " I", " help", " you", " today", "?"},
+			thinkValue:      nil,
+			expectedContent: "Hello, how can I help you today?",
+		},
 		{
 			name:             "streaming thinking then content - granular",
 			chunks:           []string{"Let", " me", " th", "ink", " about", " this", "...", "<", "/", "think", ">", "\n", "Here", " is", " my", " answer", "."},
@@ -141,6 +268,45 @@ func TestNemotron3NanoParser_Streaming(t *testing.T) {
 			expectedThinking: "Step 1: Analyze\nStep 2: Process",
 			expectedContent:  "The answer.",
 		},
+		{
+			name:       "streaming tool call - highly granular",
+			chunks:     []string{"<", "tool", "_", "call", ">", "\n", "<", "func", "tion", "=", "get", "_", "weather", ">", "\n", "<", "param", "eter", "=", "city", ">", "\n", "Par", "is", "\n", "</", "param", "eter", ">", "\n", "</", "func", "tion", ">", "\n", "</", "tool", "_", "call", ">"},
+			thinkValue: nil,
+			expectedCalls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Name:      "get_weather",
+						Arguments: testArgs(map[string]any{"city": "Paris"}),
+					},
+				},
+			},
+		},
+		{
+			name:            "streaming content then tool call - granular",
+			chunks:          []string{"Let", " me", " check", " the", " weather", ".", "\n<", "tool_call", ">", "\n", "<function=", "get_weather", ">", "\n", "<parameter=", "city", ">", "\n", "NYC", "\n", "</parameter>", "\n", "</function>", "\n", "</tool_call>"},
+			thinkValue:      nil,
+			expectedContent: "Let me check the weather.",
+			expectedCalls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Name:      "get_weather",
+						Arguments: testArgs(map[string]any{"city": "NYC"}),
+					},
+				},
+			},
+		},
+		{
+			name:   "tool call tag split character by character",
+			chunks: []string{"<", "t", "o", "o", "l", "_", "c", "a", "l", "l", ">", "\n", "<", "f", "u", "n", "c", "t", "i", "o", "n", "=", "t", "e", "s", "t", ">", "\n", "<", "/", "f", "u", "n", "c", "t", "i", "o", "n", ">", "\n", "<", "/", "t", "o", "o", "l", "_", "c", "a", "l", "l", ">"},
+			expectedCalls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Name:      "test",
+						Arguments: api.NewToolCallFunctionArguments(),
+					},
+				},
+			},
+		},
 		{
 			name:             "thinking close tag split character by character",
 			chunks:           []string{"I", "'", "m", " ", "t", "h", "i", "n", "k", "i", "n", "g", ".", ".", ".", "<", "/", "t", "h", "i", "n", "k", ">", "\n", "D", "o", "n", "e", "!"},
@@ -155,6 +321,22 @@ func TestNemotron3NanoParser_Streaming(t *testing.T) {
 			expectedThinking: "Thinking...",
 			expectedContent:  "Content here.",
 		},
+		{
+			name:       "tool call with multiple parameters - streaming",
+			chunks:     []string{"<tool_", "call>\n", "<function", "=book_", "flight>", "\n<para", "meter=", "from>\n", "SFO\n", "</param", "eter>", "\n<param", "eter=to", ">\nNYC", "\n</para", "meter>", "\n</func", "tion>\n", "</tool_", "call>"},
+			thinkValue: nil,
+			expectedCalls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Name: "book_flight",
+						Arguments: testArgs(map[string]any{
+							"from": "SFO",
+							"to":   "NYC",
+						}),
+					},
+				},
+			},
+		},
 		{
 			name:             "thinking then content then tool call - streaming",
 			chunks:           []string{"Ana", "lyzing", " your", " request", "...", "</", "think", ">\n", "I'll", " check", " that", " for", " you", ".", "\n", "<tool", "_call", ">\n", "<function", "=search", ">\n", "<parameter", "=query", ">\n", "test", " query", "\n</", "parameter", ">\n", "</function", ">\n", "</tool", "_call", ">"},
@@ -170,6 +352,45 @@ func TestNemotron3NanoParser_Streaming(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "multiple tool calls - streaming",
+			chunks: []string{
+				"<tool_call>", "\n", "<function=", "get_weather>", "\n",
+				"<parameter=", "city>\n", "San Fran", "cisco\n", "</parameter>", "\n",
+				"</function>", "\n", "</tool_call>", "\n",
+				"<tool_", "call>\n", "<function", "=get_weather", ">\n",
+				"<param", "eter=city", ">\nNew", " York\n", "</parameter>\n",
+				"</function>\n", "</tool_call>",
+			},
+			thinkValue: nil,
+			expectedCalls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Name:      "get_weather",
+						Arguments: testArgs(map[string]any{"city": "San Francisco"}),
+					},
+				},
+				{
+					Function: api.ToolCallFunction{
+						Name:      "get_weather",
+						Arguments: testArgs(map[string]any{"city": "New York"}),
+					},
+				},
+			},
+		},
+		{
+			name:       "tool call with multiline parameter - streaming",
+			chunks:     []string{"<tool_call>\n", "<function=", "create_note>\n", "<parameter=", "content>\n", "Line 1", "\nLine", " 2\n", "Line 3", "\n</parameter>\n", "</function>\n", "</tool_call>"},
+			thinkValue: nil,
+			expectedCalls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Name:      "create_note",
+						Arguments: testArgs(map[string]any{"content": "Line 1\nLine 2\nLine 3"}),
+					},
+				},
+			},
+		},
 		{
 			name:             "empty thinking block",
 			chunks:           []string{"</think>", "\n", "Just content."},
@@ -177,6 +398,12 @@ func TestNemotron3NanoParser_Streaming(t *testing.T) {
 			expectedThinking: "",
 			expectedContent:  "Just content.",
 		},
+		{
+			name:            "empty input chunks interspersed",
+			chunks:          []string{"Hello", "", " ", "", "world", "", "!"},
+			thinkValue:      nil,
+			expectedContent: "Hello world!",
+		},
 		{
 			name:             "tool call immediately after think close - no content",
 			chunks:           []string{"Analyzing...", "</think>", "\n", "<tool_call>", "\n<function=test>\n</function>\n", "</tool_call>"},
@@ -191,6 +418,25 @@ func TestNemotron3NanoParser_Streaming(t *testing.T) {
 				},
 			},
 		},
+		{
+			name:       "tool call with empty parameter value",
+			chunks:     []string{"<tool_call>\n<function=test>\n<parameter=name>\n", "\n</parameter>\n</function>\n</tool_call>"},
+			thinkValue: nil,
+			expectedCalls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Name:      "test",
+						Arguments: testArgs(map[string]any{"name": ""}),
+					},
+				},
+			},
+		},
+		{
+			name:            "partial tool call tag at end - buffered",
+			chunks:          []string{"Here's some content", "<tool"},
+			thinkValue:      nil,
+			expectedContent: "Here's some content",
+		},
 	}
 
 	for _, tt := range tests {
@@ -326,65 +572,3 @@ func TestNemotron3NanoParser_WithTools(t *testing.T) {
 		t.Errorf("calls mismatch (-got +want):\n%s", diff)
 	}
 }
-
-// TestNemotron3NanoParser_ToolCallWithoutThinkClose tests the case where thinking is enabled
-// but the model outputs content + tool call WITHOUT the </think> tag.
-// The parser should still parse the tool call (content before is treated as thinking).
-func TestNemotron3NanoParser_ToolCallWithoutThinkClose(t *testing.T) {
-	chunks := []string{
-		"Let", " me", " analyze", " this", ".", "\n",
-		"<tool_call>", "\n",
-		"<function=get_weather>", "\n",
-		"<parameter=city>", "Paris", "</parameter>", "\n",
-		"</function>", "\n",
-		"</tool_call>",
-	}
-
-	p := &Nemotron3NanoParser{}
-	p.Init(nil, nil, &api.ThinkValue{Value: true}) // thinking ENABLED but model doesn't output </think>
-
-	var allContent string
-	var allThinking string
-	var allCalls []api.ToolCall
-
-	for _, chunk := range chunks {
-		content, thinking, calls, err := p.Add(chunk, false)
-		if err != nil {
-			t.Fatalf("unexpected error: %v", err)
-		}
-		allContent += content
-		allThinking += thinking
-		allCalls = append(allCalls, calls...)
-	}
-
-	// Drain
-	content, thinking, calls, err := p.Add("", true)
-	if err != nil {
-		t.Fatalf("unexpected error on done: %v", err)
-	}
-	allContent += content
-	allThinking += thinking
-	allCalls = append(allCalls, calls...)
-
-	// The parser was in thinking mode, so text before <tool_call> is emitted as thinking.
-	expectedThinking := "Let me analyze this."
-
-	expectedCalls := []api.ToolCall{
-		{
-			Function: api.ToolCallFunction{
-				Name:      "get_weather",
-				Arguments: testArgs(map[string]any{"city": "Paris"}),
-			},
-		},
-	}
-
-	if allContent != "" {
-		t.Errorf("expected no content (text was streamed as thinking), got: %q", allContent)
-	}
-	if diff := cmp.Diff(allThinking, expectedThinking); diff != "" {
-		t.Errorf("thinking mismatch (-got +want):\n%s", diff)
-	}
-	if diff := cmp.Diff(allCalls, expectedCalls, argsComparer); diff != "" {
-		t.Errorf("calls mismatch (-got +want):\n%s", diff)
-	}
-}

model/parsers/qwen3coder_test.go

@@ -91,37 +91,6 @@ func TestQwenParserStreaming(t *testing.T) {
 				},
 			},
 		},
-		{
-			desc: "tool call tags split character by character",
-			steps: []step{
-				{input: "<", wantEvents: []qwenEvent{}},
-				{input: "t", wantEvents: []qwenEvent{}},
-				{input: "o", wantEvents: []qwenEvent{}},
-				{input: "o", wantEvents: []qwenEvent{}},
-				{input: "l", wantEvents: []qwenEvent{}},
-				{input: "_", wantEvents: []qwenEvent{}},
-				{input: "c", wantEvents: []qwenEvent{}},
-				{input: "a", wantEvents: []qwenEvent{}},
-				{input: "l", wantEvents: []qwenEvent{}},
-				{input: "l", wantEvents: []qwenEvent{}},
-				{input: ">", wantEvents: []qwenEvent{}},
-				{input: "a", wantEvents: []qwenEvent{}},
-				{input: "b", wantEvents: []qwenEvent{}},
-				{input: "c", wantEvents: []qwenEvent{}},
-				{input: "<", wantEvents: []qwenEvent{}},
-				{input: "/", wantEvents: []qwenEvent{}},
-				{input: "t", wantEvents: []qwenEvent{}},
-				{input: "o", wantEvents: []qwenEvent{}},
-				{input: "o", wantEvents: []qwenEvent{}},
-				{input: "l", wantEvents: []qwenEvent{}},
-				{input: "_", wantEvents: []qwenEvent{}},
-				{input: "c", wantEvents: []qwenEvent{}},
-				{input: "a", wantEvents: []qwenEvent{}},
-				{input: "l", wantEvents: []qwenEvent{}},
-				{input: "l", wantEvents: []qwenEvent{}},
-				{input: ">", wantEvents: []qwenEvent{qwenEventRawToolCall{raw: "abc"}}},
-			},
-		},
 		{
 			desc: "trailing whitespace between content and tool call",
 			steps: []step{

openai/openai.go

@@ -630,10 +630,6 @@ func nameFromToolCallID(messages []Message, toolCallID string) string {
 
 // decodeImageURL decodes a base64 data URI into raw image bytes.
 func decodeImageURL(url string) (api.ImageData, error) {
-	if strings.HasPrefix(url, "http://") || strings.HasPrefix(url, "https://") {
-		return nil, errors.New("image URLs are not currently supported, please use base64 encoded data instead")
-	}
-
 	types := []string{"jpeg", "jpg", "png", "webp"}
 
 	// Support blank mime type to match /api/chat's behavior of taking just unadorned base64
@@ -737,60 +733,3 @@ func FromCompleteRequest(r CompletionRequest) (api.GenerateRequest, error) {
 		DebugRenderOnly: r.DebugRenderOnly,
 	}, nil
 }
-
-// ImageGenerationRequest is an OpenAI-compatible image generation request.
-type ImageGenerationRequest struct {
-	Model          string `json:"model"`
-	Prompt         string `json:"prompt"`
-	N              int    `json:"n,omitempty"`
-	Size           string `json:"size,omitempty"`
-	ResponseFormat string `json:"response_format,omitempty"`
-	Seed           *int64 `json:"seed,omitempty"`
-}
-
-// ImageGenerationResponse is an OpenAI-compatible image generation response.
-type ImageGenerationResponse struct {
-	Created int64            `json:"created"`
-	Data    []ImageURLOrData `json:"data"`
-}
-
-// ImageURLOrData contains either a URL or base64-encoded image data.
-type ImageURLOrData struct {
-	URL     string `json:"url,omitempty"`
-	B64JSON string `json:"b64_json,omitempty"`
-}
-
-// FromImageGenerationRequest converts an OpenAI image generation request to an Ollama GenerateRequest.
-func FromImageGenerationRequest(r ImageGenerationRequest) api.GenerateRequest {
-	req := api.GenerateRequest{
-		Model:  r.Model,
-		Prompt: r.Prompt,
-	}
-	// Parse size if provided (e.g., "1024x768")
-	if r.Size != "" {
-		var w, h int32
-		if _, err := fmt.Sscanf(r.Size, "%dx%d", &w, &h); err == nil {
-			req.Width = w
-			req.Height = h
-		}
-	}
-	if r.Seed != nil {
-		if req.Options == nil {
-			req.Options = map[string]any{}
-		}
-		req.Options["seed"] = *r.Seed
-	}
-	return req
-}
-
-// ToImageGenerationResponse converts an Ollama GenerateResponse to an OpenAI ImageGenerationResponse.
-func ToImageGenerationResponse(resp api.GenerateResponse) ImageGenerationResponse {
-	var data []ImageURLOrData
-	if resp.Image != "" {
-		data = []ImageURLOrData{{B64JSON: resp.Image}}
-	}
-	return ImageGenerationResponse{
-		Created: resp.CreatedAt.Unix(),
-		Data:    data,
-	}
-}

openai/responses.go

@@ -4,7 +4,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"math/rand"
-	"time"
 
 	"github.com/ollama/ollama/api"
 )
@@ -266,9 +265,9 @@ type ResponsesText struct {
 type ResponsesTool struct {
 	Type        string         `json:"type"` // "function"
 	Name        string         `json:"name"`
-	Description *string        `json:"description"` // nullable but required
-	Strict      *bool          `json:"strict"`      // nullable but required
-	Parameters  map[string]any `json:"parameters"`  // nullable but required
+	Description string         `json:"description,omitempty"`
+	Strict      bool           `json:"strict,omitempty"`
+	Parameters  map[string]any `json:"parameters,omitempty"`
 }
 
 type ResponsesRequest struct {
@@ -476,16 +475,11 @@ func convertTool(t ResponsesTool) (api.Tool, error) {
 		}
 	}
 
-	var description string
-	if t.Description != nil {
-		description = *t.Description
-	}
-
 	return api.Tool{
 		Type: t.Type,
 		Function: api.ToolFunction{
 			Name:        t.Name,
-			Description: description,
+			Description: t.Description,
 			Parameters:  params,
 		},
 	}, nil
@@ -522,60 +516,17 @@ func convertInputMessage(m ResponsesInputMessage) (api.Message, error) {
 
 // Response types for the Responses API
 
-// ResponsesTextField represents the text output configuration in the response.
-type ResponsesTextField struct {
-	Format ResponsesTextFormat `json:"format"`
-}
-
-// ResponsesReasoningOutput represents reasoning configuration in the response.
-type ResponsesReasoningOutput struct {
-	Effort  *string `json:"effort,omitempty"`
-	Summary *string `json:"summary,omitempty"`
-}
-
-// ResponsesError represents an error in the response.
-type ResponsesError struct {
-	Code    string `json:"code"`
-	Message string `json:"message"`
-}
-
-// ResponsesIncompleteDetails represents details about why a response was incomplete.
-type ResponsesIncompleteDetails struct {
-	Reason string `json:"reason"`
-}
-
 type ResponsesResponse struct {
-	ID                 string                      `json:"id"`
-	Object             string                      `json:"object"`
-	CreatedAt          int64                       `json:"created_at"`
-	CompletedAt        *int64                      `json:"completed_at"`
-	Status             string                      `json:"status"`
-	IncompleteDetails  *ResponsesIncompleteDetails `json:"incomplete_details"`
-	Model              string                      `json:"model"`
-	PreviousResponseID *string                     `json:"previous_response_id"`
-	Instructions       *string                     `json:"instructions"`
-	Output             []ResponsesOutputItem       `json:"output"`
-	Error              *ResponsesError             `json:"error"`
-	Tools              []ResponsesTool             `json:"tools"`
-	ToolChoice         any                         `json:"tool_choice"`
-	Truncation         string                      `json:"truncation"`
-	ParallelToolCalls  bool                        `json:"parallel_tool_calls"`
-	Text               ResponsesTextField          `json:"text"`
-	TopP               float64                     `json:"top_p"`
-	PresencePenalty    float64                     `json:"presence_penalty"`
-	FrequencyPenalty   float64                     `json:"frequency_penalty"`
-	TopLogprobs        int                         `json:"top_logprobs"`
-	Temperature        float64                     `json:"temperature"`
-	Reasoning          *ResponsesReasoningOutput   `json:"reasoning"`
-	Usage              *ResponsesUsage             `json:"usage"`
-	MaxOutputTokens    *int                        `json:"max_output_tokens"`
-	MaxToolCalls       *int                        `json:"max_tool_calls"`
-	Store              bool                        `json:"store"`
-	Background         bool                        `json:"background"`
-	ServiceTier        string                      `json:"service_tier"`
-	Metadata           map[string]any              `json:"metadata"`
-	SafetyIdentifier   *string                     `json:"safety_identifier"`
-	PromptCacheKey     *string                     `json:"prompt_cache_key"`
+	ID        string                `json:"id"`
+	Object    string                `json:"object"`
+	CreatedAt int64                 `json:"created_at"`
+	Status    string                `json:"status"`
+	Model     string                `json:"model"`
+	Output    []ResponsesOutputItem `json:"output"`
+	Usage     *ResponsesUsage       `json:"usage,omitempty"`
+	// TODO(drifkin): add `temperature` and `top_p` to the response, but this
+	// requires additional plumbing to find the effective values since the
+	// defaults can come from the model or the request
 }
 
 type ResponsesOutputItem struct {
@@ -599,39 +550,18 @@ type ResponsesReasoningSummary struct {
 }
 
 type ResponsesOutputContent struct {
-	Type        string `json:"type"` // "output_text"
-	Text        string `json:"text"`
-	Annotations []any  `json:"annotations"`
-	Logprobs    []any  `json:"logprobs"`
-}
-
-type ResponsesInputTokensDetails struct {
-	CachedTokens int `json:"cached_tokens"`
-}
-
-type ResponsesOutputTokensDetails struct {
-	ReasoningTokens int `json:"reasoning_tokens"`
+	Type string `json:"type"` // "output_text"
+	Text string `json:"text"`
 }
 
 type ResponsesUsage struct {
-	InputTokens         int                          `json:"input_tokens"`
-	OutputTokens        int                          `json:"output_tokens"`
-	TotalTokens         int                          `json:"total_tokens"`
-	InputTokensDetails  ResponsesInputTokensDetails  `json:"input_tokens_details"`
-	OutputTokensDetails ResponsesOutputTokensDetails `json:"output_tokens_details"`
+	InputTokens  int `json:"input_tokens"`
+	OutputTokens int `json:"output_tokens"`
+	TotalTokens  int `json:"total_tokens"`
 }
 
-// derefFloat64 returns the value of a float64 pointer, or a default if nil.
-func derefFloat64(p *float64, def float64) float64 {
-	if p != nil {
-		return *p
-	}
-	return def
-}
-
-// ToResponse converts an api.ChatResponse to a Responses API response.
-// The request is used to echo back request parameters in the response.
-func ToResponse(model, responseID, itemID string, chatResponse api.ChatResponse, request ResponsesRequest) ResponsesResponse {
+// ToResponse converts an api.ChatResponse to a Responses API response
+func ToResponse(model, responseID, itemID string, chatResponse api.ChatResponse) ResponsesResponse {
 	var output []ResponsesOutputItem
 
 	// Add reasoning item if thinking is present
@@ -655,7 +585,6 @@ func ToResponse(model, responseID, itemID string, chatResponse api.ChatResponse,
 			output = append(output, ResponsesOutputItem{
 				ID:        fmt.Sprintf("fc_%s_%d", responseID, i),
 				Type:      "function_call",
-				Status:    "completed",
 				CallID:    tc.ID,
 				Name:      tc.Function.Name,
 				Arguments: tc.Function.Arguments,
@@ -669,90 +598,25 @@ func ToResponse(model, responseID, itemID string, chatResponse api.ChatResponse,
 			Role:   "assistant",
 			Content: []ResponsesOutputContent{
 				{
-					Type:        "output_text",
-					Text:        chatResponse.Message.Content,
-					Annotations: []any{},
-					Logprobs:    []any{},
+					Type: "output_text",
+					Text: chatResponse.Message.Content,
 				},
 			},
 		})
 	}
 
-	var instructions *string
-	if request.Instructions != "" {
-		instructions = &request.Instructions
-	}
-
-	// Build truncation with default
-	truncation := "disabled"
-	if request.Truncation != nil {
-		truncation = *request.Truncation
-	}
-
-	tools := request.Tools
-	if tools == nil {
-		tools = []ResponsesTool{}
-	}
-
-	text := ResponsesTextField{
-		Format: ResponsesTextFormat{Type: "text"},
-	}
-	if request.Text != nil && request.Text.Format != nil {
-		text.Format = *request.Text.Format
-	}
-
-	// Build reasoning output from request
-	var reasoning *ResponsesReasoningOutput
-	if request.Reasoning.Effort != "" || request.Reasoning.Summary != "" {
-		reasoning = &ResponsesReasoningOutput{}
-		if request.Reasoning.Effort != "" {
-			reasoning.Effort = &request.Reasoning.Effort
-		}
-		if request.Reasoning.Summary != "" {
-			reasoning.Summary = &request.Reasoning.Summary
-		}
-	}
-
 	return ResponsesResponse{
-		ID:                 responseID,
-		Object:             "response",
-		CreatedAt:          chatResponse.CreatedAt.Unix(),
-		CompletedAt:        nil, // Set by middleware when writing final response
-		Status:             "completed",
-		IncompleteDetails:  nil, // Only populated if response incomplete
-		Model:              model,
-		PreviousResponseID: nil, // Not supported
-		Instructions:       instructions,
-		Output:             output,
-		Error:              nil, // Only populated on failure
-		Tools:              tools,
-		ToolChoice:         "auto", // Default value
-		Truncation:         truncation,
-		ParallelToolCalls:  true, // Default value
-		Text:               text,
-		TopP:               derefFloat64(request.TopP, 1.0),
-		PresencePenalty:    0, // Default value
-		FrequencyPenalty:   0, // Default value
-		TopLogprobs:        0, // Default value
-		Temperature:        derefFloat64(request.Temperature, 1.0),
-		Reasoning:          reasoning,
+		ID:        responseID,
+		Object:    "response",
+		CreatedAt: chatResponse.CreatedAt.Unix(),
+		Status:    "completed",
+		Model:     model,
+		Output:    output,
 		Usage: &ResponsesUsage{
 			InputTokens:  chatResponse.PromptEvalCount,
 			OutputTokens: chatResponse.EvalCount,
 			TotalTokens:  chatResponse.PromptEvalCount + chatResponse.EvalCount,
-			// TODO(drifkin): wire through the actual values
-			InputTokensDetails: ResponsesInputTokensDetails{CachedTokens: 0},
-			// TODO(drifkin): wire through the actual values
-			OutputTokensDetails: ResponsesOutputTokensDetails{ReasoningTokens: 0},
 		},
-		MaxOutputTokens:  request.MaxOutputTokens,
-		MaxToolCalls:     nil,   // Not supported
-		Store:            false, // We don't store responses
-		Background:       request.Background,
-		ServiceTier:      "default", // Default value
-		Metadata:         map[string]any{},
-		SafetyIdentifier: nil, // Not supported
-		PromptCacheKey:   nil, // Not supported
 	}
 }
 
@@ -772,7 +636,6 @@ type ResponsesStreamConverter struct {
 	responseID string
 	itemID     string
 	model      string
-	request    ResponsesRequest
 
 	// State tracking (mutated across Process calls)
 	firstWrite      bool
@@ -805,12 +668,11 @@ func (c *ResponsesStreamConverter) newEvent(eventType string, data map[string]an
 }
 
 // NewResponsesStreamConverter creates a new converter with the given configuration.
-func NewResponsesStreamConverter(responseID, itemID, model string, request ResponsesRequest) *ResponsesStreamConverter {
+func NewResponsesStreamConverter(responseID, itemID, model string) *ResponsesStreamConverter {
 	return &ResponsesStreamConverter{
 		responseID: responseID,
 		itemID:     itemID,
 		model:      model,
-		request:    request,
 		firstWrite: true,
 	}
 }
@@ -855,120 +717,25 @@ func (c *ResponsesStreamConverter) Process(r api.ChatResponse) []ResponsesStream
 	return events
 }
 
-// buildResponseObject creates a full response object with all required fields for streaming events.
-func (c *ResponsesStreamConverter) buildResponseObject(status string, output []any, usage map[string]any) map[string]any {
-	var instructions any = nil
-	if c.request.Instructions != "" {
-		instructions = c.request.Instructions
-	}
-
-	truncation := "disabled"
-	if c.request.Truncation != nil {
-		truncation = *c.request.Truncation
-	}
-
-	var tools []any
-	if c.request.Tools != nil {
-		for _, t := range c.request.Tools {
-			tools = append(tools, map[string]any{
-				"type":        t.Type,
-				"name":        t.Name,
-				"description": t.Description,
-				"strict":      t.Strict,
-				"parameters":  t.Parameters,
-			})
-		}
-	}
-	if tools == nil {
-		tools = []any{}
-	}
-
-	textFormat := map[string]any{"type": "text"}
-	if c.request.Text != nil && c.request.Text.Format != nil {
-		textFormat = map[string]any{
-			"type": c.request.Text.Format.Type,
-		}
-		if c.request.Text.Format.Name != "" {
-			textFormat["name"] = c.request.Text.Format.Name
-		}
-		if c.request.Text.Format.Schema != nil {
-			textFormat["schema"] = c.request.Text.Format.Schema
-		}
-		if c.request.Text.Format.Strict != nil {
-			textFormat["strict"] = *c.request.Text.Format.Strict
-		}
-	}
-
-	var reasoning any = nil
-	if c.request.Reasoning.Effort != "" || c.request.Reasoning.Summary != "" {
-		r := map[string]any{}
-		if c.request.Reasoning.Effort != "" {
-			r["effort"] = c.request.Reasoning.Effort
-		} else {
-			r["effort"] = nil
-		}
-		if c.request.Reasoning.Summary != "" {
-			r["summary"] = c.request.Reasoning.Summary
-		} else {
-			r["summary"] = nil
-		}
-		reasoning = r
-	}
-
-	// Build top_p and temperature with defaults
-	topP := 1.0
-	if c.request.TopP != nil {
-		topP = *c.request.TopP
-	}
-	temperature := 1.0
-	if c.request.Temperature != nil {
-		temperature = *c.request.Temperature
-	}
-
-	return map[string]any{
-		"id":                   c.responseID,
-		"object":               "response",
-		"created_at":           time.Now().Unix(),
-		"completed_at":         nil,
-		"status":               status,
-		"incomplete_details":   nil,
-		"model":                c.model,
-		"previous_response_id": nil,
-		"instructions":         instructions,
-		"output":               output,
-		"error":                nil,
-		"tools":                tools,
-		"tool_choice":          "auto",
-		"truncation":           truncation,
-		"parallel_tool_calls":  true,
-		"text":                 map[string]any{"format": textFormat},
-		"top_p":                topP,
-		"presence_penalty":     0,
-		"frequency_penalty":    0,
-		"top_logprobs":         0,
-		"temperature":          temperature,
-		"reasoning":            reasoning,
-		"usage":                usage,
-		"max_output_tokens":    c.request.MaxOutputTokens,
-		"max_tool_calls":       nil,
-		"store":                false,
-		"background":           c.request.Background,
-		"service_tier":         "default",
-		"metadata":             map[string]any{},
-		"safety_identifier":    nil,
-		"prompt_cache_key":     nil,
-	}
-}
-
 func (c *ResponsesStreamConverter) createResponseCreatedEvent() ResponsesStreamEvent {
 	return c.newEvent("response.created", map[string]any{
-		"response": c.buildResponseObject("in_progress", []any{}, nil),
+		"response": map[string]any{
+			"id":     c.responseID,
+			"object": "response",
+			"status": "in_progress",
+			"output": []any{},
+		},
 	})
 }
 
 func (c *ResponsesStreamConverter) createResponseInProgressEvent() ResponsesStreamEvent {
 	return c.newEvent("response.in_progress", map[string]any{
-		"response": c.buildResponseObject("in_progress", []any{}, nil),
+		"response": map[string]any{
+			"id":     c.responseID,
+			"object": "response",
+			"status": "in_progress",
+			"output": []any{},
+		},
 	})
 }
 
@@ -995,10 +762,9 @@ func (c *ResponsesStreamConverter) processThinking(thinking string) []ResponsesS
 
 	// Emit delta
 	events = append(events, c.newEvent("response.reasoning_summary_text.delta", map[string]any{
-		"item_id":       c.reasoningItemID,
-		"output_index":  c.outputIndex,
-		"summary_index": 0,
-		"delta":         thinking,
+		"item_id":      c.reasoningItemID,
+		"output_index": c.outputIndex,
+		"delta":        thinking,
 	}))
 
 	// TODO(drifkin): consider adding
@@ -1017,10 +783,9 @@ func (c *ResponsesStreamConverter) finishReasoning() []ResponsesStreamEvent {
 
 	events := []ResponsesStreamEvent{
 		c.newEvent("response.reasoning_summary_text.done", map[string]any{
-			"item_id":       c.reasoningItemID,
-			"output_index":  c.outputIndex,
-			"summary_index": 0,
-			"text":          c.accumulatedThinking,
+			"item_id":      c.reasoningItemID,
+			"output_index": c.outputIndex,
+			"text":         c.accumulatedThinking,
 		}),
 		c.newEvent("response.output_item.done", map[string]any{
 			"output_index": c.outputIndex,
@@ -1133,10 +898,8 @@ func (c *ResponsesStreamConverter) processTextContent(content string) []Response
 			"output_index":  c.outputIndex,
 			"content_index": c.contentIndex,
 			"part": map[string]any{
-				"type":        "output_text",
-				"text":        "",
-				"annotations": []any{},
-				"logprobs":    []any{},
+				"type": "output_text",
+				"text": "",
 			},
 		}))
 	}
@@ -1150,7 +913,6 @@ func (c *ResponsesStreamConverter) processTextContent(content string) []Response
 		"output_index":  c.outputIndex,
 		"content_index": 0,
 		"delta":         content,
-		"logprobs":      []any{},
 	}))
 
 	return events
@@ -1182,10 +944,8 @@ func (c *ResponsesStreamConverter) buildFinalOutput() []any {
 			"status": "completed",
 			"role":   "assistant",
 			"content": []map[string]any{{
-				"type":        "output_text",
-				"text":        c.accumulatedText,
-				"annotations": []any{},
-				"logprobs":    []any{},
+				"type": "output_text",
+				"text": c.accumulatedText,
 			}},
 		})
 	}
@@ -1207,7 +967,6 @@ func (c *ResponsesStreamConverter) processCompletion(r api.ChatResponse) []Respo
 			"output_index":  c.outputIndex,
 			"content_index": 0,
 			"text":          c.accumulatedText,
-			"logprobs":      []any{},
 		}))
 
 		// response.content_part.done
@@ -1216,10 +975,8 @@ func (c *ResponsesStreamConverter) processCompletion(r api.ChatResponse) []Respo
 			"output_index":  c.outputIndex,
 			"content_index": 0,
 			"part": map[string]any{
-				"type":        "output_text",
-				"text":        c.accumulatedText,
-				"annotations": []any{},
-				"logprobs":    []any{},
+				"type": "output_text",
+				"text": c.accumulatedText,
 			},
 		}))
 
@@ -1232,31 +989,26 @@ func (c *ResponsesStreamConverter) processCompletion(r api.ChatResponse) []Respo
 				"status": "completed",
 				"role":   "assistant",
 				"content": []map[string]any{{
-					"type":        "output_text",
-					"text":        c.accumulatedText,
-					"annotations": []any{},
-					"logprobs":    []any{},
+					"type": "output_text",
+					"text": c.accumulatedText,
 				}},
 			},
 		}))
 	}
 
 	// response.completed
-	usage := map[string]any{
-		"input_tokens":  r.PromptEvalCount,
-		"output_tokens": r.EvalCount,
-		"total_tokens":  r.PromptEvalCount + r.EvalCount,
-		"input_tokens_details": map[string]any{
-			"cached_tokens": 0,
-		},
-		"output_tokens_details": map[string]any{
-			"reasoning_tokens": 0,
-		},
-	}
-	response := c.buildResponseObject("completed", c.buildFinalOutput(), usage)
-	response["completed_at"] = time.Now().Unix()
 	events = append(events, c.newEvent("response.completed", map[string]any{
-		"response": response,
+		"response": map[string]any{
+			"id":     c.responseID,
+			"object": "response",
+			"status": "completed",
+			"output": c.buildFinalOutput(),
+			"usage": map[string]any{
+				"input_tokens":  r.PromptEvalCount,
+				"output_tokens": r.EvalCount,
+				"total_tokens":  r.PromptEvalCount + r.EvalCount,
+			},
+		},
 	}))
 
 	return events

openai/responses_test.go

@@ -850,7 +850,7 @@ func TestFromResponsesRequest_Images(t *testing.T) {
 }
 
 func TestResponsesStreamConverter_TextOnly(t *testing.T) {
-	converter := NewResponsesStreamConverter("resp_123", "msg_456", "gpt-oss:20b", ResponsesRequest{})
+	converter := NewResponsesStreamConverter("resp_123", "msg_456", "gpt-oss:20b")
 
 	// First chunk with content
 	events := converter.Process(api.ChatResponse{
@@ -916,7 +916,7 @@ func TestResponsesStreamConverter_TextOnly(t *testing.T) {
 }
 
 func TestResponsesStreamConverter_ToolCalls(t *testing.T) {
-	converter := NewResponsesStreamConverter("resp_123", "msg_456", "gpt-oss:20b", ResponsesRequest{})
+	converter := NewResponsesStreamConverter("resp_123", "msg_456", "gpt-oss:20b")
 
 	events := converter.Process(api.ChatResponse{
 		Message: api.Message{
@@ -952,7 +952,7 @@ func TestResponsesStreamConverter_ToolCalls(t *testing.T) {
 }
 
 func TestResponsesStreamConverter_Reasoning(t *testing.T) {
-	converter := NewResponsesStreamConverter("resp_123", "msg_456", "gpt-oss:20b", ResponsesRequest{})
+	converter := NewResponsesStreamConverter("resp_123", "msg_456", "gpt-oss:20b")
 
 	// First chunk with thinking
 	events := converter.Process(api.ChatResponse{
@@ -1267,7 +1267,7 @@ func TestToResponse_WithReasoning(t *testing.T) {
 			Content:  "The answer is 42",
 		},
 		Done: true,
-	}, ResponsesRequest{})
+	})
 
 	// Should have 2 output items: reasoning + message
 	if len(response.Output) != 2 {
@@ -1638,7 +1638,7 @@ func TestFromResponsesRequest_ShorthandFormats(t *testing.T) {
 
 func TestResponsesStreamConverter_OutputIncludesContent(t *testing.T) {
 	// Verify that response.output_item.done includes content field for messages
-	converter := NewResponsesStreamConverter("resp_123", "msg_456", "gpt-oss:20b", ResponsesRequest{})
+	converter := NewResponsesStreamConverter("resp_123", "msg_456", "gpt-oss:20b")
 
 	// First chunk
 	converter.Process(api.ChatResponse{
@@ -1686,7 +1686,7 @@ func TestResponsesStreamConverter_OutputIncludesContent(t *testing.T) {
 
 func TestResponsesStreamConverter_ResponseCompletedIncludesOutput(t *testing.T) {
 	// Verify that response.completed includes the output array
-	converter := NewResponsesStreamConverter("resp_123", "msg_456", "gpt-oss:20b", ResponsesRequest{})
+	converter := NewResponsesStreamConverter("resp_123", "msg_456", "gpt-oss:20b")
 
 	// Process some content
 	converter.Process(api.ChatResponse{
@@ -1730,7 +1730,7 @@ func TestResponsesStreamConverter_ResponseCompletedIncludesOutput(t *testing.T)
 
 func TestResponsesStreamConverter_ResponseCreatedIncludesOutput(t *testing.T) {
 	// Verify that response.created includes an empty output array
-	converter := NewResponsesStreamConverter("resp_123", "msg_456", "gpt-oss:20b", ResponsesRequest{})
+	converter := NewResponsesStreamConverter("resp_123", "msg_456", "gpt-oss:20b")
 
 	events := converter.Process(api.ChatResponse{
 		Message: api.Message{Content: "Hi"},
@@ -1757,7 +1757,7 @@ func TestResponsesStreamConverter_ResponseCreatedIncludesOutput(t *testing.T) {
 
 func TestResponsesStreamConverter_SequenceNumbers(t *testing.T) {
 	// Verify that events include incrementing sequence numbers
-	converter := NewResponsesStreamConverter("resp_123", "msg_456", "gpt-oss:20b", ResponsesRequest{})
+	converter := NewResponsesStreamConverter("resp_123", "msg_456", "gpt-oss:20b")
 
 	events := converter.Process(api.ChatResponse{
 		Message: api.Message{Content: "Hello"},
@@ -1791,7 +1791,7 @@ func TestResponsesStreamConverter_SequenceNumbers(t *testing.T) {
 
 func TestResponsesStreamConverter_FunctionCallStatus(t *testing.T) {
 	// Verify that function call items include status field
-	converter := NewResponsesStreamConverter("resp_123", "msg_456", "gpt-oss:20b", ResponsesRequest{})
+	converter := NewResponsesStreamConverter("resp_123", "msg_456", "gpt-oss:20b")
 
 	events := converter.Process(api.ChatResponse{
 		Message: api.Message{

readline/readline.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"io"
 	"os"
-	"strings"
 )
 
 type Prompt struct {
@@ -37,11 +36,10 @@ type Terminal struct {
 }
 
 type Instance struct {
-	Prompt      *Prompt
-	Terminal    *Terminal
-	History     *History
-	Pasting     bool
-	pastedLines []string
+	Prompt   *Prompt
+	Terminal *Terminal
+	History  *History
+	Pasting  bool
 }
 
 func New(prompt Prompt) (*Instance, error) {
@@ -176,8 +174,6 @@ func (i *Instance) Readline() (string, error) {
 		case CharEsc:
 			esc = true
 		case CharInterrupt:
-			i.pastedLines = nil
-			i.Prompt.UseAlt = false
 			return "", ErrInterrupt
 		case CharPrev:
 			i.historyPrev(buf, &currentLineBuf)
@@ -192,23 +188,7 @@ func (i *Instance) Readline() (string, error) {
 		case CharForward:
 			buf.MoveRight()
 		case CharBackspace, CharCtrlH:
-			if buf.IsEmpty() && len(i.pastedLines) > 0 {
-				lastIdx := len(i.pastedLines) - 1
-				prevLine := i.pastedLines[lastIdx]
-				i.pastedLines = i.pastedLines[:lastIdx]
-				fmt.Print(CursorBOL + ClearToEOL + CursorUp + CursorBOL + ClearToEOL)
-				if len(i.pastedLines) == 0 {
-					fmt.Print(i.Prompt.Prompt)
-					i.Prompt.UseAlt = false
-				} else {
-					fmt.Print(i.Prompt.AltPrompt)
-				}
-				for _, r := range prevLine {
-					buf.Add(r)
-				}
-			} else {
-				buf.Remove()
-			}
+			buf.Remove()
 		case CharTab:
 			// todo: convert back to real tabs
 			for range 8 {
@@ -231,28 +211,13 @@ func (i *Instance) Readline() (string, error) {
 		case CharCtrlZ:
 			fd := os.Stdin.Fd()
 			return handleCharCtrlZ(fd, i.Terminal.termios)
-		case CharCtrlJ:
-			i.pastedLines = append(i.pastedLines, buf.String())
-			buf.Buf.Clear()
-			buf.Pos = 0
-			buf.DisplayPos = 0
-			buf.LineHasSpace.Clear()
-			fmt.Println()
-			fmt.Print(i.Prompt.AltPrompt)
-			i.Prompt.UseAlt = true
-			continue
-		case CharEnter:
+		case CharEnter, CharCtrlJ:
 			output := buf.String()
-			if len(i.pastedLines) > 0 {
-				output = strings.Join(i.pastedLines, "\n") + "\n" + output
-				i.pastedLines = nil
-			}
 			if output != "" {
 				i.History.Add(output)
 			}
 			buf.MoveToEnd()
 			fmt.Println()
-			i.Prompt.UseAlt = false
 
 			return output, nil
 		default:

runner/runner.go

@@ -7,13 +7,26 @@ import (
 )
 
 func Execute(args []string) error {
-	if len(args) > 0 {
-		switch args[0] {
-		case "--ollama-engine":
-			return ollamarunner.Execute(args[1:])
-		case "--image-engine":
-			return imagerunner.Execute(args[1:])
-		}
+	if args[0] == "runner" {
+		args = args[1:]
+	}
+
+	var newRunner bool
+	var imageRunner bool
+	if len(args) > 0 && args[0] == "--ollama-engine" {
+		args = args[1:]
+		newRunner = true
+	}
+	if len(args) > 0 && args[0] == "--image-engine" {
+		args = args[1:]
+		imageRunner = true
+	}
+
+	if imageRunner {
+		return imagerunner.Execute(args)
+	} else if newRunner {
+		return ollamarunner.Execute(args)
+	} else {
+		return llamarunner.Execute(args)
 	}
-	return llamarunner.Execute(args)
 }

scripts/build_darwin.sh

@@ -60,7 +60,7 @@ _build_darwin() {
             cmake --install $BUILD_DIR --component MLX
             # Override CGO flags to point to the amd64 build directory
             MLX_CGO_CFLAGS="-O3 -I$(pwd)/$BUILD_DIR/_deps/mlx-c-src -mmacosx-version-min=14.0"
-            MLX_CGO_LDFLAGS="-ldl -lc++ -framework Accelerate -mmacosx-version-min=14.0"
+            MLX_CGO_LDFLAGS="-L$(pwd)/$BUILD_DIR/lib/ollama -lmlxc -lmlx -Wl,-rpath,@executable_path -lc++ -framework Accelerate -mmacosx-version-min=14.0"
         else
             BUILD_DIR=build
             cmake --preset MLX \
@@ -71,12 +71,10 @@ _build_darwin() {
             cmake --install $BUILD_DIR --component MLX
             # Use default CGO flags from mlx.go for arm64
             MLX_CGO_CFLAGS="-O3 -I$(pwd)/$BUILD_DIR/_deps/mlx-c-src -mmacosx-version-min=14.0"
-            MLX_CGO_LDFLAGS="-lc++ -framework Metal -framework Foundation -framework Accelerate -mmacosx-version-min=14.0"
+            MLX_CGO_LDFLAGS="-L$(pwd)/$BUILD_DIR/lib/ollama -lmlxc -lmlx -Wl,-rpath,@executable_path -lc++ -framework Metal -framework Foundation -framework Accelerate -mmacosx-version-min=14.0"
         fi
-        GOOS=darwin GOARCH=$ARCH CGO_ENABLED=1 CGO_CFLAGS="$MLX_CGO_CFLAGS" CGO_LDFLAGS="$MLX_CGO_LDFLAGS" go build -tags mlx -o $INSTALL_PREFIX .
-        # Copy MLX libraries to same directory as executable for dlopen
-        cp $INSTALL_PREFIX/lib/ollama/libmlxc.dylib $INSTALL_PREFIX/
-        cp $INSTALL_PREFIX/lib/ollama/libmlx.dylib $INSTALL_PREFIX/
+        GOOS=darwin GOARCH=$ARCH CGO_ENABLED=1 CGO_CFLAGS="$MLX_CGO_CFLAGS" CGO_LDFLAGS="$MLX_CGO_LDFLAGS" go build -tags mlx -o $INSTALL_PREFIX/ollama-mlx .
+        GOOS=darwin GOARCH=$ARCH CGO_ENABLED=1 go build -o $INSTALL_PREFIX .
     done
 }
 
@@ -84,10 +82,12 @@ _sign_darwin() {
     status "Creating universal binary..."
     mkdir -p dist/darwin
     lipo -create -output dist/darwin/ollama dist/darwin-*/ollama
+    lipo -create -output dist/darwin/ollama-mlx dist/darwin-*/ollama-mlx
     chmod +x dist/darwin/ollama
+    chmod +x dist/darwin/ollama-mlx
 
     if [ -n "$APPLE_IDENTITY" ]; then
-        for F in dist/darwin/ollama dist/darwin-*/lib/ollama/*; do
+        for F in dist/darwin/ollama dist/darwin-*/lib/ollama/* dist/darwin/ollama-mlx; do
             codesign -f --timestamp -s "$APPLE_IDENTITY" --identifier ai.ollama.ollama --options=runtime $F
         done
 
@@ -154,6 +154,7 @@ _build_macapp() {
     mkdir -p dist/Ollama.app/Contents/Resources
     if [ -d dist/darwin-amd64 ]; then
         lipo -create -output dist/Ollama.app/Contents/Resources/ollama dist/darwin-amd64/ollama dist/darwin-arm64/ollama
+        lipo -create -output dist/Ollama.app/Contents/Resources/ollama-mlx dist/darwin-amd64/ollama-mlx dist/darwin-arm64/ollama-mlx
         for F in dist/darwin-amd64/lib/ollama/*mlx*.dylib ; do
             lipo -create -output dist/darwin/$(basename $F) $F dist/darwin-arm64/lib/ollama/$(basename $F)
         done
@@ -165,27 +166,28 @@ _build_macapp() {
         cp -a dist/darwin/ollama dist/Ollama.app/Contents/Resources/ollama
         cp dist/darwin/*.so dist/darwin/*.dylib dist/Ollama.app/Contents/Resources/
     fi
+    cp -a dist/darwin/ollama-mlx dist/Ollama.app/Contents/Resources/ollama-mlx
     chmod a+x dist/Ollama.app/Contents/Resources/ollama
 
     # Sign
     if [ -n "$APPLE_IDENTITY" ]; then
         codesign -f --timestamp -s "$APPLE_IDENTITY" --identifier ai.ollama.ollama --options=runtime dist/Ollama.app/Contents/Resources/ollama
-        for lib in dist/Ollama.app/Contents/Resources/*.so dist/Ollama.app/Contents/Resources/*.dylib dist/Ollama.app/Contents/Resources/*.metallib ; do
+        for lib in dist/Ollama.app/Contents/Resources/*.so dist/Ollama.app/Contents/Resources/*.dylib dist/Ollama.app/Contents/Resources/*.metallib dist/Ollama.app/Contents/Resources/ollama-mlx ; do
             codesign -f --timestamp -s "$APPLE_IDENTITY" --identifier ai.ollama.ollama --options=runtime ${lib}
         done
         codesign -f --timestamp -s "$APPLE_IDENTITY" --identifier com.electron.ollama --deep --options=runtime dist/Ollama.app
     fi
 
     rm -f dist/Ollama-darwin.zip
-    ditto -c -k --norsrc --keepParent dist/Ollama.app dist/Ollama-darwin.zip
-    (cd dist/Ollama.app/Contents/Resources/; tar -cf - ollama *.so *.dylib *.metallib 2>/dev/null) | gzip -9vc > dist/ollama-darwin.tgz
+    ditto -c -k --keepParent dist/Ollama.app dist/Ollama-darwin.zip
+    (cd dist/Ollama.app/Contents/Resources/; tar -cf - ollama ollama-mlx *.so *.dylib *.metallib 2>/dev/null) | gzip -9vc > dist/ollama-darwin.tgz
 
     # Notarize and Staple
     if [ -n "$APPLE_IDENTITY" ]; then
         $(xcrun -f notarytool) submit dist/Ollama-darwin.zip --wait --timeout 20m --apple-id "$APPLE_ID" --password "$APPLE_PASSWORD" --team-id "$APPLE_TEAM_ID"
         rm -f dist/Ollama-darwin.zip
         $(xcrun -f stapler) staple dist/Ollama.app
-        ditto -c -k --norsrc --keepParent dist/Ollama.app dist/Ollama-darwin.zip
+        ditto -c -k --keepParent dist/Ollama.app dist/Ollama-darwin.zip
 
         rm -f dist/Ollama.dmg
 

server/auth.go

@@ -50,17 +50,12 @@ func (r registryChallenge) URL() (*url.URL, error) {
 	return redirectURL, nil
 }
 
-func getAuthorizationToken(ctx context.Context, challenge registryChallenge, originalHost string) (string, error) {
+func getAuthorizationToken(ctx context.Context, challenge registryChallenge) (string, error) {
 	redirectURL, err := challenge.URL()
 	if err != nil {
 		return "", err
 	}
 
-	// Validate that the realm host matches the original request host to prevent sending tokens cross-origin.
-	if redirectURL.Host != originalHost {
-		return "", fmt.Errorf("realm host %q does not match original host %q", redirectURL.Host, originalHost)
-	}
-
 	sha256sum := sha256.Sum256(nil)
 	data := []byte(fmt.Sprintf("%s,%s,%s", http.MethodGet, redirectURL.String(), base64.StdEncoding.EncodeToString([]byte(hex.EncodeToString(sha256sum[:])))))
 

server/auth_test.go

@@ -1,113 +0,0 @@
-package server
-
-import (
-	"context"
-	"strings"
-	"testing"
-	"time"
-)
-
-func TestGetAuthorizationTokenRejectsCrossDomain(t *testing.T) {
-	tests := []struct {
-		realm        string
-		originalHost string
-		wantMismatch bool
-	}{
-		{"https://example.com/token", "example.com", false},
-		{"https://example.com/token", "other.com", true},
-		{"https://example.com/token", "localhost:8000", true},
-		{"https://localhost:5000/token", "localhost:5000", false},
-		{"https://localhost:5000/token", "localhost:6000", true},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.originalHost, func(t *testing.T) {
-			ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
-			defer cancel()
-
-			challenge := registryChallenge{Realm: tt.realm, Service: "test", Scope: "repo:x:pull"}
-			_, err := getAuthorizationToken(ctx, challenge, tt.originalHost)
-
-			isMismatch := err != nil && strings.Contains(err.Error(), "does not match")
-			if tt.wantMismatch && !isMismatch {
-				t.Errorf("expected domain mismatch error, got: %v", err)
-			}
-			if !tt.wantMismatch && isMismatch {
-				t.Errorf("unexpected domain mismatch error: %v", err)
-			}
-		})
-	}
-}
-
-func TestParseRegistryChallenge(t *testing.T) {
-	tests := []struct {
-		input                             string
-		wantRealm, wantService, wantScope string
-	}{
-		{
-			`Bearer realm="https://auth.example.com/token",service="registry",scope="repo:foo:pull"`,
-			"https://auth.example.com/token", "registry", "repo:foo:pull",
-		},
-		{
-			`Bearer realm="https://r.ollama.ai/v2/token",service="ollama",scope="-"`,
-			"https://r.ollama.ai/v2/token", "ollama", "-",
-		},
-		{"", "", "", ""},
-	}
-
-	for _, tt := range tests {
-		result := parseRegistryChallenge(tt.input)
-		if result.Realm != tt.wantRealm || result.Service != tt.wantService || result.Scope != tt.wantScope {
-			t.Errorf("parseRegistryChallenge(%q) = {%q, %q, %q}, want {%q, %q, %q}",
-				tt.input, result.Realm, result.Service, result.Scope,
-				tt.wantRealm, tt.wantService, tt.wantScope)
-		}
-	}
-}
-
-func TestRegistryChallengeURL(t *testing.T) {
-	challenge := registryChallenge{
-		Realm:   "https://auth.example.com/token",
-		Service: "registry",
-		Scope:   "repo:foo:pull repo:bar:push",
-	}
-
-	u, err := challenge.URL()
-	if err != nil {
-		t.Fatalf("URL() error: %v", err)
-	}
-
-	if u.Host != "auth.example.com" {
-		t.Errorf("host = %q, want %q", u.Host, "auth.example.com")
-	}
-	if u.Path != "/token" {
-		t.Errorf("path = %q, want %q", u.Path, "/token")
-	}
-
-	q := u.Query()
-	if q.Get("service") != "registry" {
-		t.Errorf("service = %q, want %q", q.Get("service"), "registry")
-	}
-	if scopes := q["scope"]; len(scopes) != 2 {
-		t.Errorf("scope count = %d, want 2", len(scopes))
-	}
-	if q.Get("ts") == "" {
-		t.Error("missing ts")
-	}
-	if q.Get("nonce") == "" {
-		t.Error("missing nonce")
-	}
-
-	// Nonces should differ between calls
-	u2, _ := challenge.URL()
-	if q.Get("nonce") == u2.Query().Get("nonce") {
-		t.Error("nonce should be unique per call")
-	}
-}
-
-func TestRegistryChallengeURLInvalid(t *testing.T) {
-	challenge := registryChallenge{Realm: "://invalid"}
-	if _, err := challenge.URL(); err == nil {
-		t.Error("expected error for invalid URL")
-	}
-}

server/download.go

@@ -95,11 +95,48 @@ func (p *blobDownloadPart) UnmarshalJSON(b []byte) error {
 }
 
 const (
-	numDownloadParts          = 16
+	// numDownloadParts is the default number of concurrent download parts for standard downloads
+	numDownloadParts = 16
+	// numHFDownloadParts is the reduced number of concurrent download parts for HuggingFace
+	// downloads to avoid triggering rate limits (HTTP 429 errors). See GitHub issue #13297.
+	numHFDownloadParts        = 4
 	minDownloadPartSize int64 = 100 * format.MegaByte
 	maxDownloadPartSize int64 = 1000 * format.MegaByte
 )
 
+// isHuggingFaceURL returns true if the URL is from a HuggingFace domain.
+// This includes:
+//   - huggingface.co (main domain)
+//   - *.huggingface.co (subdomains like cdn-lfs.huggingface.co)
+//   - hf.co (shortlink domain)
+//   - *.hf.co (CDN domains like cdn-lfs.hf.co, cdn-lfs3.hf.co)
+func isHuggingFaceURL(u *url.URL) bool {
+	if u == nil {
+		return false
+	}
+	host := strings.ToLower(u.Hostname())
+	return host == "huggingface.co" ||
+		strings.HasSuffix(host, ".huggingface.co") ||
+		host == "hf.co" ||
+		strings.HasSuffix(host, ".hf.co")
+}
+
+// getNumDownloadParts returns the number of concurrent download parts to use
+// for the given URL. HuggingFace URLs use reduced concurrency (default 4) to
+// avoid triggering rate limits. This can be overridden via the OLLAMA_HF_CONCURRENCY
+// environment variable. For non-HuggingFace URLs, returns the standard concurrency (16).
+func getNumDownloadParts(u *url.URL) int {
+	if isHuggingFaceURL(u) {
+		if v := os.Getenv("OLLAMA_HF_CONCURRENCY"); v != "" {
+			if n, err := strconv.Atoi(v); err == nil && n > 0 {
+				return n
+			}
+		}
+		return numHFDownloadParts
+	}
+	return numDownloadParts
+}
+
 func (p *blobDownloadPart) Name() string {
 	return strings.Join([]string{
 		p.blobDownload.Name, "partial", strconv.Itoa(p.N),
@@ -271,7 +308,11 @@ func (b *blobDownload) run(ctx context.Context, requestURL *url.URL, opts *regis
 	}
 
 	g, inner := errgroup.WithContext(ctx)
-	g.SetLimit(numDownloadParts)
+	concurrency := getNumDownloadParts(directURL)
+	if concurrency != numDownloadParts {
+		slog.Info(fmt.Sprintf("using reduced concurrency (%d) for HuggingFace download", concurrency))
+	}
+	g.SetLimit(concurrency)
 	for i := range b.Parts {
 		part := b.Parts[i]
 		if part.Completed.Load() == part.Size {

server/download_test.go

@@ -0,0 +1,194 @@
+package server
+
+import (
+	"net/url"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestIsHuggingFaceURL(t *testing.T) {
+	tests := []struct {
+		name     string
+		url      string
+		expected bool
+	}{
+		{
+			name:     "nil url",
+			url:      "",
+			expected: false,
+		},
+		{
+			name:     "huggingface.co main domain",
+			url:      "https://huggingface.co/some/model",
+			expected: true,
+		},
+		{
+			name:     "cdn-lfs.huggingface.co subdomain",
+			url:      "https://cdn-lfs.huggingface.co/repos/abc/123",
+			expected: true,
+		},
+		{
+			name:     "cdn-lfs3.hf.co CDN domain",
+			url:      "https://cdn-lfs3.hf.co/repos/abc/123",
+			expected: true,
+		},
+		{
+			name:     "hf.co shortlink domain",
+			url:      "https://hf.co/model",
+			expected: true,
+		},
+		{
+			name:     "uppercase HuggingFace domain",
+			url:      "https://HUGGINGFACE.CO/model",
+			expected: true,
+		},
+		{
+			name:     "mixed case HF domain",
+			url:      "https://Cdn-Lfs.HF.Co/repos",
+			expected: true,
+		},
+		{
+			name:     "ollama registry",
+			url:      "https://registry.ollama.ai/v2/library/llama3",
+			expected: false,
+		},
+		{
+			name:     "github.com",
+			url:      "https://github.com/ollama/ollama",
+			expected: false,
+		},
+		{
+			name:     "fake huggingface domain",
+			url:      "https://nothuggingface.co/model",
+			expected: false,
+		},
+		{
+			name:     "fake hf domain",
+			url:      "https://nothf.co/model",
+			expected: false,
+		},
+		{
+			name:     "huggingface in path not host",
+			url:      "https://example.com/huggingface.co/model",
+			expected: false,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			var u *url.URL
+			if tc.url != "" {
+				var err error
+				u, err = url.Parse(tc.url)
+				if err != nil {
+					t.Fatalf("failed to parse URL: %v", err)
+				}
+			}
+			got := isHuggingFaceURL(u)
+			assert.Equal(t, tc.expected, got)
+		})
+	}
+}
+
+func TestGetNumDownloadParts(t *testing.T) {
+	tests := []struct {
+		name        string
+		url         string
+		envValue    string
+		expected    int
+		description string
+	}{
+		{
+			name:        "nil url returns default",
+			url:         "",
+			envValue:    "",
+			expected:    numDownloadParts,
+			description: "nil URL should return standard concurrency",
+		},
+		{
+			name:        "ollama registry returns default",
+			url:         "https://registry.ollama.ai/v2/library/llama3",
+			envValue:    "",
+			expected:    numDownloadParts,
+			description: "Ollama registry should use standard concurrency",
+		},
+		{
+			name:        "huggingface returns reduced default",
+			url:         "https://huggingface.co/model/repo",
+			envValue:    "",
+			expected:    numHFDownloadParts,
+			description: "HuggingFace should use reduced concurrency",
+		},
+		{
+			name:        "hf.co CDN returns reduced default",
+			url:         "https://cdn-lfs3.hf.co/repos/abc/123",
+			envValue:    "",
+			expected:    numHFDownloadParts,
+			description: "HuggingFace CDN should use reduced concurrency",
+		},
+		{
+			name:        "huggingface with env override",
+			url:         "https://huggingface.co/model/repo",
+			envValue:    "2",
+			expected:    2,
+			description: "OLLAMA_HF_CONCURRENCY should override default",
+		},
+		{
+			name:        "huggingface with higher env override",
+			url:         "https://huggingface.co/model/repo",
+			envValue:    "8",
+			expected:    8,
+			description: "OLLAMA_HF_CONCURRENCY can be set higher than default",
+		},
+		{
+			name:        "huggingface with invalid env (non-numeric)",
+			url:         "https://huggingface.co/model/repo",
+			envValue:    "invalid",
+			expected:    numHFDownloadParts,
+			description: "Invalid OLLAMA_HF_CONCURRENCY should fall back to default",
+		},
+		{
+			name:        "huggingface with invalid env (zero)",
+			url:         "https://huggingface.co/model/repo",
+			envValue:    "0",
+			expected:    numHFDownloadParts,
+			description: "Zero OLLAMA_HF_CONCURRENCY should fall back to default",
+		},
+		{
+			name:        "huggingface with invalid env (negative)",
+			url:         "https://huggingface.co/model/repo",
+			envValue:    "-1",
+			expected:    numHFDownloadParts,
+			description: "Negative OLLAMA_HF_CONCURRENCY should fall back to default",
+		},
+		{
+			name:        "non-huggingface ignores env",
+			url:         "https://registry.ollama.ai/v2/library/llama3",
+			envValue:    "2",
+			expected:    numDownloadParts,
+			description: "OLLAMA_HF_CONCURRENCY should not affect non-HF URLs",
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			// Set or clear the environment variable
+			if tc.envValue != "" {
+				t.Setenv("OLLAMA_HF_CONCURRENCY", tc.envValue)
+			}
+
+			var u *url.URL
+			if tc.url != "" {
+				var err error
+				u, err = url.Parse(tc.url)
+				if err != nil {
+					t.Fatalf("failed to parse URL: %v", err)
+				}
+			}
+
+			got := getNumDownloadParts(u)
+			assert.Equal(t, tc.expected, got, tc.description)
+		})
+	}
+}

server/images.go

@@ -41,7 +41,6 @@ var (
 	errCapabilityVision     = errors.New("vision")
 	errCapabilityEmbedding  = errors.New("embedding")
 	errCapabilityThinking   = errors.New("thinking")
-	errCapabilityImage      = errors.New("image generation")
 	errInsecureProtocol     = errors.New("insecure protocol http")
 )
 
@@ -77,7 +76,7 @@ func (m *Model) Capabilities() []model.Capability {
 
 	// Check for image generation model via config capabilities
 	if slices.Contains(m.Config.Capabilities, "image") {
-		return []model.Capability{model.CapabilityImage}
+		return []model.Capability{model.CapabilityImageGeneration}
 	}
 
 	// Check for completion capability
@@ -160,7 +159,6 @@ func (m *Model) CheckCapabilities(want ...model.Capability) error {
 		model.CapabilityVision:     errCapabilityVision,
 		model.CapabilityEmbedding:  errCapabilityEmbedding,
 		model.CapabilityThinking:   errCapabilityThinking,
-		model.CapabilityImage:      errCapabilityImage,
 	}
 
 	for _, cap := range want {
@@ -777,7 +775,7 @@ func pullWithTransfer(ctx context.Context, mp ModelPath, layers []Layer, manifes
 			Realm:   challenge.Realm,
 			Service: challenge.Service,
 			Scope:   challenge.Scope,
-		}, base.Host)
+		})
 	}
 
 	if err := transfer.Download(ctx, transfer.DownloadOptions{
@@ -852,7 +850,7 @@ func pushWithTransfer(ctx context.Context, mp ModelPath, layers []Layer, manifes
 			Realm:   challenge.Realm,
 			Service: challenge.Service,
 			Scope:   challenge.Scope,
-		}, base.Host)
+		})
 	}
 
 	return transfer.Upload(ctx, transfer.UploadOptions{
@@ -918,7 +916,7 @@ func makeRequestWithRetry(ctx context.Context, method string, requestURL *url.UR
 
 			// Handle authentication error with one retry
 			challenge := parseRegistryChallenge(resp.Header.Get("www-authenticate"))
-			token, err := getAuthorizationToken(ctx, challenge, requestURL.Host)
+			token, err := getAuthorizationToken(ctx, challenge)
 			if err != nil {
 				return nil, err
 			}

server/images_test.go

@@ -54,7 +54,7 @@ func TestModelCapabilities(t *testing.T) {
 					Capabilities: []string{"image"},
 				},
 			},
-			expectedCaps: []model.Capability{model.CapabilityImage},
+			expectedCaps: []model.Capability{model.CapabilityImageGeneration},
 		},
 		{
 			name: "model with completion capability",
@@ -242,24 +242,6 @@ func TestModelCheckCapabilities(t *testing.T) {
 			checkCaps:      []model.Capability{"unknown"},
 			expectedErrMsg: "unknown capability",
 		},
-		{
-			name: "model missing image generation capability",
-			model: Model{
-				ModelPath: completionModelPath,
-				Template:  chatTemplate,
-			},
-			checkCaps:      []model.Capability{model.CapabilityImage},
-			expectedErrMsg: "does not support image generation",
-		},
-		{
-			name: "model with image generation capability",
-			model: Model{
-				Config: model.ConfigV2{
-					Capabilities: []string{"image"},
-				},
-			},
-			checkCaps: []model.Capability{model.CapabilityImage},
-		},
 	}
 
 	for _, tt := range tests {

server/routes.go

@@ -51,7 +51,7 @@ import (
 	"github.com/ollama/ollama/types/model"
 	"github.com/ollama/ollama/version"
 	"github.com/ollama/ollama/x/imagegen"
-	xserver "github.com/ollama/ollama/x/server"
+	imagegenapi "github.com/ollama/ollama/x/imagegen/api"
 )
 
 const signinURLStr = "https://ollama.com/connect?name=%s&key=%s"
@@ -164,6 +164,29 @@ func (s *Server) scheduleRunner(ctx context.Context, name string, caps []model.C
 	return runner.llama, model, &opts, nil
 }
 
+// ScheduleImageGenRunner schedules an image generation model runner.
+// This implements the imagegenapi.RunnerScheduler interface.
+func (s *Server) ScheduleImageGenRunner(c *gin.Context, modelName string, opts api.Options, keepAlive *api.Duration) (llm.LlamaServer, error) {
+	m := &Model{
+		Name:      modelName,
+		ShortName: modelName,
+		ModelPath: modelName, // For image gen, ModelPath is just the model name
+		Config: model.ConfigV2{
+			Capabilities: []string{"image"},
+		},
+	}
+
+	runnerCh, errCh := s.sched.GetRunner(c.Request.Context(), m, opts, keepAlive)
+	var runner *runnerRef
+	select {
+	case runner = <-runnerCh:
+	case err := <-errCh:
+		return nil, err
+	}
+
+	return runner.llama, nil
+}
+
 func signinURL() (string, error) {
 	pubKey, err := auth.GetPublicKey()
 	if err != nil {
@@ -191,6 +214,12 @@ func (s *Server) GenerateHandler(c *gin.Context) {
 		return
 	}
 
+	// Check if this is a known image generation model
+	if imagegen.ResolveModelName(req.Model) != "" {
+		imagegenapi.HandleGenerateRequest(c, s, req.Model, req.Prompt, req.KeepAlive, streamResponse)
+		return
+	}
+
 	name := model.ParseName(req.Model)
 	if !name.IsValid() {
 		// Ideally this is "invalid model name" but we're keeping with
@@ -220,12 +249,6 @@ func (s *Server) GenerateHandler(c *gin.Context) {
 		return
 	}
 
-	// Handle image generation models
-	if slices.Contains(m.Capabilities(), model.CapabilityImage) {
-		s.handleImageGenerate(c, req, name.String(), checkpointStart)
-		return
-	}
-
 	if req.TopLogprobs < 0 || req.TopLogprobs > 20 {
 		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "top_logprobs must be between 0 and 20"})
 		return
@@ -1102,7 +1125,7 @@ func GetModelInfo(req api.ShowRequest) (*api.ShowResponse, error) {
 	}
 
 	// For image generation models, populate details from imagegen package
-	if slices.Contains(m.Capabilities(), model.CapabilityImage) {
+	if slices.Contains(m.Capabilities(), model.CapabilityImageGeneration) {
 		if info, err := imagegen.GetModelInfo(name.String()); err == nil {
 			modelDetails.Family = info.Architecture
 			modelDetails.ParameterSize = format.HumanNumber(uint64(info.ParameterCount))
@@ -1110,22 +1133,6 @@ func GetModelInfo(req api.ShowRequest) (*api.ShowResponse, error) {
 		}
 	}
 
-	// For safetensors LLM models (experimental), populate details from config.json
-	if m.Config.ModelFormat == "safetensors" && slices.Contains(m.Config.Capabilities, "completion") {
-		if info, err := xserver.GetSafetensorsLLMInfo(name.String()); err == nil {
-			if arch, ok := info["general.architecture"].(string); ok && arch != "" {
-				modelDetails.Family = arch
-			}
-			if paramCount, ok := info["general.parameter_count"].(int64); ok && paramCount > 0 {
-				modelDetails.ParameterSize = format.HumanNumber(uint64(paramCount))
-			}
-		}
-		// Get torch_dtype directly from config.json for quantization level
-		if dtype, err := xserver.GetSafetensorsDtype(name.String()); err == nil && dtype != "" {
-			modelDetails.QuantizationLevel = dtype
-		}
-	}
-
 	if req.System != "" {
 		m.System = req.System
 	}
@@ -1208,27 +1215,7 @@ func GetModelInfo(req api.ShowRequest) (*api.ShowResponse, error) {
 		return resp, nil
 	}
 
-	if slices.Contains(m.Capabilities(), model.CapabilityImage) {
-		// Populate tensor info if verbose
-		if req.Verbose {
-			if tensors, err := xserver.GetSafetensorsTensorInfo(name.String()); err == nil {
-				resp.Tensors = tensors
-			}
-		}
-		return resp, nil
-	}
-
-	// For safetensors LLM models (experimental), populate ModelInfo from config.json
-	if m.Config.ModelFormat == "safetensors" && slices.Contains(m.Config.Capabilities, "completion") {
-		if info, err := xserver.GetSafetensorsLLMInfo(name.String()); err == nil {
-			resp.ModelInfo = info
-		}
-		// Populate tensor info if verbose
-		if req.Verbose {
-			if tensors, err := xserver.GetSafetensorsTensorInfo(name.String()); err == nil {
-				resp.Tensors = tensors
-			}
-		}
+	if slices.Contains(m.Capabilities(), model.CapabilityImageGeneration) {
 		return resp, nil
 	}
 
@@ -1600,12 +1587,13 @@ func (s *Server) GenerateRoutes(rc *ollama.Registry) (http.Handler, error) {
 	r.GET("/v1/models", middleware.ListMiddleware(), s.ListHandler)
 	r.GET("/v1/models/:model", middleware.RetrieveMiddleware(), s.ShowHandler)
 	r.POST("/v1/responses", middleware.ResponsesMiddleware(), s.ChatHandler)
-	// OpenAI-compatible image generation endpoint
-	r.POST("/v1/images/generations", middleware.ImageGenerationsMiddleware(), s.GenerateHandler)
 
 	// Inference (Anthropic compatibility)
 	r.POST("/v1/messages", middleware.AnthropicMessagesMiddleware(), s.ChatHandler)
 
+	// Experimental image generation support
+	imagegenapi.RegisterRoutes(r, s)
+
 	if rc != nil {
 		// wrap old with new
 		rs := &registry.Local{
@@ -2472,91 +2460,3 @@ func filterThinkTags(msgs []api.Message, m *Model) []api.Message {
 	}
 	return msgs
 }
-
-// handleImageGenerate handles image generation requests within GenerateHandler.
-// This is called when the model has the Image capability.
-func (s *Server) handleImageGenerate(c *gin.Context, req api.GenerateRequest, modelName string, checkpointStart time.Time) {
-	// Validate image dimensions
-	const maxDimension int32 = 4096
-	if req.Width > maxDimension || req.Height > maxDimension {
-		c.JSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("width and height must be <= %d", maxDimension)})
-		return
-	}
-
-	// Schedule the runner for image generation
-	runner, _, _, err := s.scheduleRunner(c.Request.Context(), modelName, []model.Capability{model.CapabilityImage}, nil, req.KeepAlive)
-	if err != nil {
-		handleScheduleError(c, req.Model, err)
-		return
-	}
-
-	checkpointLoaded := time.Now()
-
-	// Handle load-only request (empty prompt)
-	if req.Prompt == "" {
-		c.JSON(http.StatusOK, api.GenerateResponse{
-			Model:      req.Model,
-			CreatedAt:  time.Now().UTC(),
-			Done:       true,
-			DoneReason: "load",
-		})
-		return
-	}
-
-	// Set headers for streaming response
-	c.Header("Content-Type", "application/x-ndjson")
-
-	// Get seed from options if provided
-	var seed int64
-	if s, ok := req.Options["seed"]; ok {
-		switch v := s.(type) {
-		case int:
-			seed = int64(v)
-		case int64:
-			seed = v
-		case float64:
-			seed = int64(v)
-		}
-	}
-
-	var streamStarted bool
-	if err := runner.Completion(c.Request.Context(), llm.CompletionRequest{
-		Prompt: req.Prompt,
-		Width:  req.Width,
-		Height: req.Height,
-		Steps:  req.Steps,
-		Seed:   seed,
-	}, func(cr llm.CompletionResponse) {
-		streamStarted = true
-		res := api.GenerateResponse{
-			Model:     req.Model,
-			CreatedAt: time.Now().UTC(),
-			Done:      cr.Done,
-		}
-
-		if cr.TotalSteps > 0 {
-			res.Completed = int64(cr.Step)
-			res.Total = int64(cr.TotalSteps)
-		}
-
-		if cr.Image != "" {
-			res.Image = cr.Image
-		}
-
-		if cr.Done {
-			res.DoneReason = cr.DoneReason.String()
-			res.Metrics.TotalDuration = time.Since(checkpointStart)
-			res.Metrics.LoadDuration = checkpointLoaded.Sub(checkpointStart)
-		}
-
-		data, _ := json.Marshal(res)
-		c.Writer.Write(append(data, '\n'))
-		c.Writer.Flush()
-	}); err != nil {
-		// Only send JSON error if streaming hasn't started yet
-		// (once streaming starts, headers are committed and we can't change status code)
-		if !streamStarted {
-			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
-		}
-	}
-}

server/sched.go

@@ -574,8 +574,7 @@ func (s *Scheduler) loadImageGen(req *LlmRequest) bool {
 		Options:         &req.opts,
 		loading:         false,
 		sessionDuration: sessionDuration,
-		totalSize:       server.TotalSize(),
-		vramSize:        server.VRAMSize(),
+		refCount:        1,
 	}
 
 	s.loadedMu.Lock()

server/sched_test.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"log/slog"
 	"os"
+	"slices"
 	"testing"
 	"time"
 
@@ -16,6 +17,7 @@ import (
 	"github.com/ollama/ollama/fs/ggml"
 	"github.com/ollama/ollama/llm"
 	"github.com/ollama/ollama/ml"
+	"github.com/ollama/ollama/types/model"
 )
 
 func TestMain(m *testing.M) {
@@ -805,8 +807,32 @@ func (s *mockLlm) GetDeviceInfos(ctx context.Context) []ml.DeviceInfo { return n
 func (s *mockLlm) HasExited() bool                                    { return false }
 func (s *mockLlm) GetActiveDeviceIDs() []ml.DeviceID                  { return nil }
 
+// TestImageGenCapabilityDetection verifies that models with "image" capability
+// are correctly identified and routed differently from language models.
+func TestImageGenCapabilityDetection(t *testing.T) {
+	// Model with image capability should be detected
+	imageModel := &Model{
+		Config: model.ConfigV2{
+			Capabilities: []string{"image"},
+		},
+	}
+	require.True(t, slices.Contains(imageModel.Config.Capabilities, "image"))
+
+	// Model without image capability should not be detected
+	langModel := &Model{
+		Config: model.ConfigV2{
+			Capabilities: []string{"completion"},
+		},
+	}
+	require.False(t, slices.Contains(langModel.Config.Capabilities, "image"))
+
+	// Empty capabilities should not match
+	emptyModel := &Model{}
+	require.False(t, slices.Contains(emptyModel.Config.Capabilities, "image"))
+}
+
 // TestImageGenRunnerCanBeEvicted verifies that an image generation model
-// loaded in the scheduler can be evicted when idle.
+// loaded in the scheduler can be evicted by a language model request.
 func TestImageGenRunnerCanBeEvicted(t *testing.T) {
 	ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond)
 	defer done()
@@ -838,59 +864,3 @@ func TestImageGenRunnerCanBeEvicted(t *testing.T) {
 	require.NotNil(t, runner)
 	require.Equal(t, "/fake/image/model", runner.modelPath)
 }
-
-// TestImageGenSchedulerCoexistence verifies that image generation models
-// can coexist with language models in the scheduler and VRAM is tracked correctly.
-func TestImageGenSchedulerCoexistence(t *testing.T) {
-	ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond)
-	defer done()
-
-	s := InitScheduler(ctx)
-	s.getGpuFn = getGpuFn
-	s.getSystemInfoFn = getSystemInfoFn
-
-	// Load both an imagegen runner and a language model runner
-	imageGenRunner := &runnerRef{
-		model:           &Model{Name: "flux", ModelPath: "/fake/flux/model"},
-		modelPath:       "/fake/flux/model",
-		llama:           &mockLlm{vramSize: 8 * format.GigaByte, vramByGPU: map[ml.DeviceID]uint64{{Library: "Metal"}: 8 * format.GigaByte}},
-		sessionDuration: 10 * time.Millisecond,
-		numParallel:     1,
-		refCount:        0,
-	}
-
-	langModelRunner := &runnerRef{
-		model:           &Model{Name: "llama3", ModelPath: "/fake/llama3/model"},
-		modelPath:       "/fake/llama3/model",
-		llama:           &mockLlm{vramSize: 4 * format.GigaByte, vramByGPU: map[ml.DeviceID]uint64{{Library: "Metal"}: 4 * format.GigaByte}},
-		sessionDuration: 10 * time.Millisecond,
-		numParallel:     1,
-		refCount:        0,
-	}
-
-	s.loadedMu.Lock()
-	s.loaded["/fake/flux/model"] = imageGenRunner
-	s.loaded["/fake/llama3/model"] = langModelRunner
-	s.loadedMu.Unlock()
-
-	// Verify both are loaded
-	s.loadedMu.Lock()
-	require.Len(t, s.loaded, 2)
-	require.NotNil(t, s.loaded["/fake/flux/model"])
-	require.NotNil(t, s.loaded["/fake/llama3/model"])
-	s.loadedMu.Unlock()
-
-	// Verify updateFreeSpace accounts for both
-	gpus := []ml.DeviceInfo{
-		{
-			DeviceID:    ml.DeviceID{Library: "Metal"},
-			TotalMemory: 24 * format.GigaByte,
-			FreeMemory:  24 * format.GigaByte,
-		},
-	}
-	s.updateFreeSpace(gpus)
-
-	// Free memory should be reduced by both models
-	expectedFree := uint64(24*format.GigaByte) - uint64(8*format.GigaByte) - uint64(4*format.GigaByte)
-	require.Equal(t, expectedFree, gpus[0].FreeMemory)
-}

server/upload.go

@@ -279,7 +279,7 @@ func (b *blobUpload) uploadPart(ctx context.Context, method string, requestURL *
 	case resp.StatusCode == http.StatusUnauthorized:
 		w.Rollback()
 		challenge := parseRegistryChallenge(resp.Header.Get("www-authenticate"))
-		token, err := getAuthorizationToken(ctx, challenge, requestURL.Host)
+		token, err := getAuthorizationToken(ctx, challenge)
 		if err != nil {
 			return err
 		}

types/model/capability.go

@@ -9,7 +9,7 @@ const (
 	CapabilityVision          = Capability("vision")
 	CapabilityEmbedding       = Capability("embedding")
 	CapabilityThinking        = Capability("thinking")
-	CapabilityImage = Capability("image")
+	CapabilityImageGeneration = Capability("image")
 )
 
 func (c Capability) String() string {

x/README.md

@@ -0,0 +1,50 @@
+# Experimental Features
+
+## MLX Backend
+
+We're working on a new experimental backend based on the [MLX project](https://github.com/ml-explore/mlx)
+
+Support is currently limited to MacOS and Linux with CUDA GPUs. We're looking to add support for Windows CUDA soon, and other GPU vendors.
+
+### Building ollama-mlx
+
+The `ollama-mlx` binary is a separate build of Ollama with MLX support enabled. This enables experimental features like image generation.
+
+#### macOS (Apple Silicon and Intel)
+
+```bash
+# Build MLX backend libraries
+cmake --preset MLX
+cmake --build --preset MLX --parallel
+cmake --install build --component MLX
+
+# Build ollama-mlx binary
+go build -tags mlx -o ollama-mlx .
+```
+
+#### Linux (CUDA)
+
+On Linux, use the preset "MLX CUDA 13" or "MLX CUDA 12" to enable CUDA with the default Ollama NVIDIA GPU architectures enabled:
+
+```bash
+# Build MLX backend libraries with CUDA support
+cmake --preset 'MLX CUDA 13'
+cmake --build --preset 'MLX CUDA 13' --parallel
+cmake --install build --component MLX
+
+# Build ollama-mlx binary
+CGO_CFLAGS="-O3 -I$(pwd)/build/_deps/mlx-c-src" \
+CGO_LDFLAGS="-L$(pwd)/build/lib/ollama -lmlxc -lmlx" \
+go build -tags mlx -o ollama-mlx .
+```
+
+#### Using build scripts
+
+The build scripts automatically create the `ollama-mlx` binary:
+
+- **macOS**: `./scripts/build_darwin.sh` produces `dist/darwin/ollama-mlx`
+- **Linux**: `./scripts/build_linux.sh` produces `ollama-mlx` in the output archives
+
+## Image Generation
+
+Image generation is built into the `ollama-mlx` binary. Run `ollama-mlx serve` to start the server with image generation support enabled.

x/cmd/run.go

@@ -25,6 +25,14 @@ import (
 	"github.com/ollama/ollama/x/tools"
 )
 
+// MultilineState tracks the state of multiline input
+type MultilineState int
+
+const (
+	MultilineNone MultilineState = iota
+	MultilineSystem
+)
+
 // Tool output capping constants
 const (
 	// localModelTokenLimit is the token limit for local models (smaller context).
@@ -648,7 +656,7 @@ func GenerateInteractive(cmd *cobra.Command, modelName string, wordWrap bool, op
 		Prompt:         ">>> ",
 		AltPrompt:      "... ",
 		Placeholder:    "Send a message (/? for help)",
-		AltPlaceholder: "Press Enter to send",
+		AltPlaceholder: `Use """ to end multi-line input`,
 	})
 	if err != nil {
 		return err
@@ -699,6 +707,7 @@ func GenerateInteractive(cmd *cobra.Command, modelName string, wordWrap bool, op
 	var sb strings.Builder
 	var format string
 	var system string
+	var multiline MultilineState = MultilineNone
 
 	for {
 		line, err := scanner.Readline()
@@ -712,12 +721,37 @@ func GenerateInteractive(cmd *cobra.Command, modelName string, wordWrap bool, op
 			}
 			scanner.Prompt.UseAlt = false
 			sb.Reset()
+			multiline = MultilineNone
 			continue
 		case err != nil:
 			return err
 		}
 
 		switch {
+		case multiline != MultilineNone:
+			// check if there's a multiline terminating string
+			before, ok := strings.CutSuffix(line, `"""`)
+			sb.WriteString(before)
+			if !ok {
+				fmt.Fprintln(&sb)
+				continue
+			}
+
+			switch multiline {
+			case MultilineSystem:
+				system = sb.String()
+				newMessage := api.Message{Role: "system", Content: system}
+				if len(messages) > 0 && messages[len(messages)-1].Role == "system" {
+					messages[len(messages)-1] = newMessage
+				} else {
+					messages = append(messages, newMessage)
+				}
+				fmt.Println("Set system message.")
+				sb.Reset()
+			}
+
+			multiline = MultilineNone
+			scanner.Prompt.UseAlt = false
 		case strings.HasPrefix(line, "/exit"), strings.HasPrefix(line, "/bye"):
 			return nil
 		case strings.HasPrefix(line, "/clear"):
@@ -826,18 +860,41 @@ func GenerateInteractive(cmd *cobra.Command, modelName string, wordWrap bool, op
 					options[args[2]] = fp[args[2]]
 				case "system":
 					if len(args) < 3 {
-						fmt.Println("Usage: /set system <message>")
+						fmt.Println("Usage: /set system <message> or /set system \"\"\"<multi-line message>\"\"\"")
 						continue
 					}
 
-					system = strings.Join(args[2:], " ")
-					newMessage := api.Message{Role: "system", Content: system}
+					multiline = MultilineSystem
+
+					line := strings.Join(args[2:], " ")
+					line, ok := strings.CutPrefix(line, `"""`)
+					if !ok {
+						multiline = MultilineNone
+					} else {
+						// only cut suffix if the line is multiline
+						line, ok = strings.CutSuffix(line, `"""`)
+						if ok {
+							multiline = MultilineNone
+						}
+					}
+
+					sb.WriteString(line)
+					if multiline != MultilineNone {
+						scanner.Prompt.UseAlt = true
+						continue
+					}
+
+					system = sb.String()
+					newMessage := api.Message{Role: "system", Content: sb.String()}
+					// Check if the slice is not empty and the last message is from 'system'
 					if len(messages) > 0 && messages[len(messages)-1].Role == "system" {
+						// Replace the last message
 						messages[len(messages)-1] = newMessage
 					} else {
 						messages = append(messages, newMessage)
 					}
 					fmt.Println("Set system message.")
+					sb.Reset()
 					continue
 				default:
 					fmt.Printf("Unknown command '/set %s'. Type /? for help\n", args[1])
@@ -1024,7 +1081,7 @@ func GenerateInteractive(cmd *cobra.Command, modelName string, wordWrap bool, op
 			sb.WriteString(line)
 		}
 
-		if sb.Len() > 0 {
+		if sb.Len() > 0 && multiline == MultilineNone {
 			newMessage := api.Message{Role: "user", Content: sb.String()}
 			messages = append(messages, newMessage)
 

x/create/client/create.go

@@ -1,282 +0,0 @@
-// Package client provides client-side model creation for safetensors-based models.
-//
-// This package is in x/ because the safetensors model storage format is under development.
-// It also exists to break an import cycle: server imports x/create, so x/create
-// cannot import server. This sub-package can import server because server doesn't
-// import it.
-package client
-
-import (
-	"bytes"
-	"encoding/json"
-	"fmt"
-	"io"
-
-	"github.com/ollama/ollama/progress"
-	"github.com/ollama/ollama/server"
-	"github.com/ollama/ollama/types/model"
-	"github.com/ollama/ollama/x/create"
-)
-
-// MinOllamaVersion is the minimum Ollama version required for safetensors models.
-const MinOllamaVersion = "0.14.0"
-
-// ModelfileConfig holds configuration extracted from a Modelfile.
-type ModelfileConfig struct {
-	Template string
-	System   string
-	License  string
-}
-
-// CreateOptions holds all options for model creation.
-type CreateOptions struct {
-	ModelName string
-	ModelDir  string
-	Quantize  string           // "fp8" for quantization
-	Modelfile *ModelfileConfig // template/system/license from Modelfile
-}
-
-// CreateModel imports a model from a local directory.
-// This creates blobs and manifest directly on disk, bypassing the HTTP API.
-// Automatically detects model type (safetensors LLM vs image gen) and routes accordingly.
-func CreateModel(opts CreateOptions, p *progress.Progress) error {
-	// Detect model type
-	isSafetensors := create.IsSafetensorsModelDir(opts.ModelDir)
-	isImageGen := create.IsTensorModelDir(opts.ModelDir)
-
-	if !isSafetensors && !isImageGen {
-		return fmt.Errorf("%s is not a supported model directory (needs config.json + *.safetensors or model_index.json)", opts.ModelDir)
-	}
-
-	// Determine model type settings
-	var modelType, spinnerKey string
-	var capabilities []string
-	if isSafetensors {
-		modelType = "safetensors model"
-		spinnerKey = "create"
-		capabilities = []string{"completion"}
-	} else {
-		modelType = "image generation model"
-		spinnerKey = "imagegen"
-		capabilities = []string{"image"}
-	}
-
-	// Set up progress spinner
-	statusMsg := "importing " + modelType
-	spinner := progress.NewSpinner(statusMsg)
-	p.Add(spinnerKey, spinner)
-
-	progressFn := func(msg string) {
-		spinner.Stop()
-		statusMsg = msg
-		spinner = progress.NewSpinner(statusMsg)
-		p.Add(spinnerKey, spinner)
-	}
-
-	// Create the model using shared callbacks
-	var err error
-	if isSafetensors {
-		err = create.CreateSafetensorsModel(
-			opts.ModelName, opts.ModelDir, opts.Quantize,
-			newLayerCreator(), newTensorLayerCreator(),
-			newManifestWriter(opts, capabilities),
-			progressFn,
-		)
-	} else {
-		err = create.CreateImageGenModel(
-			opts.ModelName, opts.ModelDir, opts.Quantize,
-			newLayerCreator(), newTensorLayerCreator(),
-			newManifestWriter(opts, capabilities),
-			progressFn,
-		)
-	}
-
-	spinner.Stop()
-	if err != nil {
-		return err
-	}
-
-	fmt.Printf("Created %s '%s'\n", modelType, opts.ModelName)
-	return nil
-}
-
-// newLayerCreator returns a LayerCreator callback for creating config/JSON layers.
-func newLayerCreator() create.LayerCreator {
-	return func(r io.Reader, mediaType, name string) (create.LayerInfo, error) {
-		layer, err := server.NewLayer(r, mediaType)
-		if err != nil {
-			return create.LayerInfo{}, err
-		}
-
-		return create.LayerInfo{
-			Digest:    layer.Digest,
-			Size:      layer.Size,
-			MediaType: layer.MediaType,
-			Name:      name,
-		}, nil
-	}
-}
-
-// newTensorLayerCreator returns a QuantizingTensorLayerCreator callback for creating tensor layers.
-// When quantize is non-empty, returns multiple layers (weight + scales + optional qbias).
-func newTensorLayerCreator() create.QuantizingTensorLayerCreator {
-	return func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]create.LayerInfo, error) {
-		if quantize != "" {
-			return createQuantizedLayers(r, name, dtype, shape, quantize)
-		}
-		return createUnquantizedLayer(r, name)
-	}
-}
-
-// createQuantizedLayers quantizes a tensor and returns the resulting layers.
-func createQuantizedLayers(r io.Reader, name, dtype string, shape []int32, quantize string) ([]create.LayerInfo, error) {
-	if !QuantizeSupported() {
-		return nil, fmt.Errorf("quantization requires MLX support")
-	}
-
-	// Quantize the tensor
-	qweightData, scalesData, qbiasData, _, _, _, err := quantizeTensor(r, name, dtype, shape, quantize)
-	if err != nil {
-		return nil, fmt.Errorf("failed to quantize %s: %w", name, err)
-	}
-
-	// Create layer for quantized weight
-	weightLayer, err := server.NewLayer(bytes.NewReader(qweightData), server.MediaTypeImageTensor)
-	if err != nil {
-		return nil, err
-	}
-
-	// Create layer for scales
-	scalesLayer, err := server.NewLayer(bytes.NewReader(scalesData), server.MediaTypeImageTensor)
-	if err != nil {
-		return nil, err
-	}
-
-	layers := []create.LayerInfo{
-		{
-			Digest:    weightLayer.Digest,
-			Size:      weightLayer.Size,
-			MediaType: weightLayer.MediaType,
-			Name:      name,
-		},
-		{
-			Digest:    scalesLayer.Digest,
-			Size:      scalesLayer.Size,
-			MediaType: scalesLayer.MediaType,
-			Name:      name + "_scale",
-		},
-	}
-
-	// Add qbiases layer if present (affine mode)
-	if qbiasData != nil {
-		qbiasLayer, err := server.NewLayer(bytes.NewReader(qbiasData), server.MediaTypeImageTensor)
-		if err != nil {
-			return nil, err
-		}
-		layers = append(layers, create.LayerInfo{
-			Digest:    qbiasLayer.Digest,
-			Size:      qbiasLayer.Size,
-			MediaType: qbiasLayer.MediaType,
-			Name:      name + "_qbias",
-		})
-	}
-
-	return layers, nil
-}
-
-// createUnquantizedLayer creates a single tensor layer without quantization.
-func createUnquantizedLayer(r io.Reader, name string) ([]create.LayerInfo, error) {
-	layer, err := server.NewLayer(r, server.MediaTypeImageTensor)
-	if err != nil {
-		return nil, err
-	}
-
-	return []create.LayerInfo{
-		{
-			Digest:    layer.Digest,
-			Size:      layer.Size,
-			MediaType: layer.MediaType,
-			Name:      name,
-		},
-	}, nil
-}
-
-// newManifestWriter returns a ManifestWriter callback for writing the model manifest.
-func newManifestWriter(opts CreateOptions, capabilities []string) create.ManifestWriter {
-	return func(modelName string, config create.LayerInfo, layers []create.LayerInfo) error {
-		name := model.ParseName(modelName)
-		if !name.IsValid() {
-			return fmt.Errorf("invalid model name: %s", modelName)
-		}
-
-		// Create config blob with version requirement
-		configData := model.ConfigV2{
-			ModelFormat:  "safetensors",
-			Capabilities: capabilities,
-			Requires:     MinOllamaVersion,
-		}
-		configJSON, err := json.Marshal(configData)
-		if err != nil {
-			return fmt.Errorf("failed to marshal config: %w", err)
-		}
-
-		// Create config layer blob
-		configLayer, err := server.NewLayer(bytes.NewReader(configJSON), "application/vnd.docker.container.image.v1+json")
-		if err != nil {
-			return fmt.Errorf("failed to create config layer: %w", err)
-		}
-
-		// Convert LayerInfo to server.Layer
-		serverLayers := make([]server.Layer, 0, len(layers))
-		for _, l := range layers {
-			serverLayers = append(serverLayers, server.Layer{
-				MediaType: l.MediaType,
-				Digest:    l.Digest,
-				Size:      l.Size,
-				Name:      l.Name,
-			})
-		}
-
-		// Add Modelfile layers if present
-		if opts.Modelfile != nil {
-			modelfileLayers, err := createModelfileLayers(opts.Modelfile)
-			if err != nil {
-				return err
-			}
-			serverLayers = append(serverLayers, modelfileLayers...)
-		}
-
-		return server.WriteManifest(name, configLayer, serverLayers)
-	}
-}
-
-// createModelfileLayers creates layers for template, system, and license from Modelfile config.
-func createModelfileLayers(mf *ModelfileConfig) ([]server.Layer, error) {
-	var layers []server.Layer
-
-	if mf.Template != "" {
-		layer, err := server.NewLayer(bytes.NewReader([]byte(mf.Template)), "application/vnd.ollama.image.template")
-		if err != nil {
-			return nil, fmt.Errorf("failed to create template layer: %w", err)
-		}
-		layers = append(layers, layer)
-	}
-
-	if mf.System != "" {
-		layer, err := server.NewLayer(bytes.NewReader([]byte(mf.System)), "application/vnd.ollama.image.system")
-		if err != nil {
-			return nil, fmt.Errorf("failed to create system layer: %w", err)
-		}
-		layers = append(layers, layer)
-	}
-
-	if mf.License != "" {
-		layer, err := server.NewLayer(bytes.NewReader([]byte(mf.License)), "application/vnd.ollama.image.license")
-		if err != nil {
-			return nil, fmt.Errorf("failed to create license layer: %w", err)
-		}
-		layers = append(layers, layer)
-	}
-
-	return layers, nil
-}

x/create/client/create_test.go

@@ -1,146 +0,0 @@
-package client
-
-import (
-	"testing"
-)
-
-func TestModelfileConfig(t *testing.T) {
-	// Test that ModelfileConfig struct works as expected
-	config := &ModelfileConfig{
-		Template: "{{ .Prompt }}",
-		System:   "You are a helpful assistant.",
-		License:  "MIT",
-	}
-
-	if config.Template != "{{ .Prompt }}" {
-		t.Errorf("Template = %q, want %q", config.Template, "{{ .Prompt }}")
-	}
-	if config.System != "You are a helpful assistant." {
-		t.Errorf("System = %q, want %q", config.System, "You are a helpful assistant.")
-	}
-	if config.License != "MIT" {
-		t.Errorf("License = %q, want %q", config.License, "MIT")
-	}
-}
-
-func TestModelfileConfig_Empty(t *testing.T) {
-	config := &ModelfileConfig{}
-
-	if config.Template != "" {
-		t.Errorf("Template should be empty, got %q", config.Template)
-	}
-	if config.System != "" {
-		t.Errorf("System should be empty, got %q", config.System)
-	}
-	if config.License != "" {
-		t.Errorf("License should be empty, got %q", config.License)
-	}
-}
-
-func TestModelfileConfig_PartialFields(t *testing.T) {
-	// Test config with only some fields set
-	config := &ModelfileConfig{
-		Template: "{{ .Prompt }}",
-		// System and License intentionally empty
-	}
-
-	if config.Template == "" {
-		t.Error("Template should not be empty")
-	}
-	if config.System != "" {
-		t.Error("System should be empty")
-	}
-	if config.License != "" {
-		t.Error("License should be empty")
-	}
-}
-
-func TestMinOllamaVersion(t *testing.T) {
-	// Verify the minimum version constant is set
-	if MinOllamaVersion == "" {
-		t.Error("MinOllamaVersion should not be empty")
-	}
-	if MinOllamaVersion != "0.14.0" {
-		t.Errorf("MinOllamaVersion = %q, want %q", MinOllamaVersion, "0.14.0")
-	}
-}
-
-func TestCreateModel_InvalidDir(t *testing.T) {
-	// Test that CreateModel returns error for invalid directory
-	err := CreateModel(CreateOptions{
-		ModelName: "test-model",
-		ModelDir:  "/nonexistent/path",
-	}, nil)
-	if err == nil {
-		t.Error("expected error for nonexistent directory, got nil")
-	}
-}
-
-func TestCreateModel_NotSafetensorsDir(t *testing.T) {
-	// Test that CreateModel returns error for directory without safetensors
-	dir := t.TempDir()
-
-	err := CreateModel(CreateOptions{
-		ModelName: "test-model",
-		ModelDir:  dir,
-	}, nil)
-	if err == nil {
-		t.Error("expected error for empty directory, got nil")
-	}
-}
-
-func TestCreateOptions(t *testing.T) {
-	opts := CreateOptions{
-		ModelName: "my-model",
-		ModelDir:  "/path/to/model",
-		Quantize:  "fp8",
-		Modelfile: &ModelfileConfig{
-			Template: "test",
-			System:   "system",
-			License:  "MIT",
-		},
-	}
-
-	if opts.ModelName != "my-model" {
-		t.Errorf("ModelName = %q, want %q", opts.ModelName, "my-model")
-	}
-	if opts.ModelDir != "/path/to/model" {
-		t.Errorf("ModelDir = %q, want %q", opts.ModelDir, "/path/to/model")
-	}
-	if opts.Quantize != "fp8" {
-		t.Errorf("Quantize = %q, want %q", opts.Quantize, "fp8")
-	}
-	if opts.Modelfile == nil {
-		t.Error("Modelfile should not be nil")
-	}
-	if opts.Modelfile.Template != "test" {
-		t.Errorf("Modelfile.Template = %q, want %q", opts.Modelfile.Template, "test")
-	}
-}
-
-func TestCreateOptions_Defaults(t *testing.T) {
-	opts := CreateOptions{
-		ModelName: "test",
-		ModelDir:  "/tmp",
-	}
-
-	// Quantize should default to empty
-	if opts.Quantize != "" {
-		t.Errorf("Quantize should be empty by default, got %q", opts.Quantize)
-	}
-
-	// Modelfile should default to nil
-	if opts.Modelfile != nil {
-		t.Error("Modelfile should be nil by default")
-	}
-}
-
-func TestQuantizeSupported(t *testing.T) {
-	// This just verifies the function exists and returns a boolean
-	// The actual value depends on build tags (mlx vs non-mlx)
-	supported := QuantizeSupported()
-
-	// In non-mlx builds, this should be false
-	// We can't easily test both cases, so just verify it returns something
-	_ = supported
-}

x/create/create.go

@@ -1,399 +0,0 @@
-package create
-
-import (
-	"encoding/json"
-	"fmt"
-	"io"
-	"os"
-	"path/filepath"
-	"slices"
-	"strings"
-
-	"github.com/ollama/ollama/envconfig"
-	"github.com/ollama/ollama/x/imagegen/safetensors"
-)
-
-// ModelConfig represents the config blob stored with a model.
-type ModelConfig struct {
-	ModelFormat  string   `json:"model_format"`
-	Capabilities []string `json:"capabilities"`
-}
-
-// Manifest represents the manifest JSON structure.
-type Manifest struct {
-	SchemaVersion int             `json:"schemaVersion"`
-	MediaType     string          `json:"mediaType"`
-	Config        ManifestLayer   `json:"config"`
-	Layers        []ManifestLayer `json:"layers"`
-}
-
-// ManifestLayer represents a layer in the manifest.
-type ManifestLayer struct {
-	MediaType string `json:"mediaType"`
-	Digest    string `json:"digest"`
-	Size      int64  `json:"size"`
-	Name      string `json:"name,omitempty"`
-}
-
-// defaultManifestDir returns the manifest storage directory.
-func defaultManifestDir() string {
-	return filepath.Join(envconfig.Models(), "manifests")
-}
-
-// defaultBlobDir returns the blob storage directory.
-func defaultBlobDir() string {
-	return filepath.Join(envconfig.Models(), "blobs")
-}
-
-// resolveManifestPath converts a model name to a manifest file path.
-func resolveManifestPath(modelName string) string {
-	host := "registry.ollama.ai"
-	namespace := "library"
-	name := modelName
-	tag := "latest"
-
-	if idx := strings.LastIndex(name, ":"); idx != -1 {
-		tag = name[idx+1:]
-		name = name[:idx]
-	}
-
-	parts := strings.Split(name, "/")
-	switch len(parts) {
-	case 3:
-		host = parts[0]
-		namespace = parts[1]
-		name = parts[2]
-	case 2:
-		namespace = parts[0]
-		name = parts[1]
-	}
-
-	return filepath.Join(defaultManifestDir(), host, namespace, name, tag)
-}
-
-// loadManifest loads a manifest for the given model name.
-func loadManifest(modelName string) (*Manifest, error) {
-	manifestPath := resolveManifestPath(modelName)
-
-	data, err := os.ReadFile(manifestPath)
-	if err != nil {
-		return nil, err
-	}
-
-	var manifest Manifest
-	if err := json.Unmarshal(data, &manifest); err != nil {
-		return nil, err
-	}
-
-	return &manifest, nil
-}
-
-// loadModelConfig loads the config blob for a model.
-func loadModelConfig(modelName string) (*ModelConfig, error) {
-	manifest, err := loadManifest(modelName)
-	if err != nil {
-		return nil, err
-	}
-
-	// Read the config blob
-	blobName := strings.Replace(manifest.Config.Digest, ":", "-", 1)
-	blobPath := filepath.Join(defaultBlobDir(), blobName)
-
-	data, err := os.ReadFile(blobPath)
-	if err != nil {
-		return nil, err
-	}
-
-	var config ModelConfig
-	if err := json.Unmarshal(data, &config); err != nil {
-		return nil, err
-	}
-
-	return &config, nil
-}
-
-// IsSafetensorsModel checks if a model was created with the experimental
-// safetensors builder by checking the model format in the config.
-func IsSafetensorsModel(modelName string) bool {
-	config, err := loadModelConfig(modelName)
-	if err != nil {
-		return false
-	}
-	return config.ModelFormat == "safetensors"
-}
-
-// IsSafetensorsLLMModel checks if a model is a safetensors LLM model
-// (has completion capability, not image generation).
-func IsSafetensorsLLMModel(modelName string) bool {
-	config, err := loadModelConfig(modelName)
-	if err != nil {
-		return false
-	}
-	return config.ModelFormat == "safetensors" && slices.Contains(config.Capabilities, "completion")
-}
-
-// IsImageGenModel checks if a model is an image generation model
-// (has image capability).
-func IsImageGenModel(modelName string) bool {
-	config, err := loadModelConfig(modelName)
-	if err != nil {
-		return false
-	}
-	return config.ModelFormat == "safetensors" && slices.Contains(config.Capabilities, "image")
-}
-
-// GetModelArchitecture returns the architecture from the model's config.json layer.
-func GetModelArchitecture(modelName string) (string, error) {
-	manifest, err := loadManifest(modelName)
-	if err != nil {
-		return "", err
-	}
-
-	// Find the config.json layer
-	for _, layer := range manifest.Layers {
-		if layer.Name == "config.json" && layer.MediaType == "application/vnd.ollama.image.json" {
-			blobName := strings.Replace(layer.Digest, ":", "-", 1)
-			blobPath := filepath.Join(defaultBlobDir(), blobName)
-
-			data, err := os.ReadFile(blobPath)
-			if err != nil {
-				return "", err
-			}
-
-			var cfg struct {
-				Architectures []string `json:"architectures"`
-				ModelType     string   `json:"model_type"`
-			}
-			if err := json.Unmarshal(data, &cfg); err != nil {
-				return "", err
-			}
-
-			// Prefer model_type, fall back to first architecture
-			if cfg.ModelType != "" {
-				return cfg.ModelType, nil
-			}
-			if len(cfg.Architectures) > 0 {
-				return cfg.Architectures[0], nil
-			}
-		}
-	}
-
-	return "", fmt.Errorf("architecture not found in model config")
-}
-
-// IsTensorModelDir checks if the directory contains a diffusers-style tensor model
-// by looking for model_index.json, which is the standard diffusers pipeline config.
-func IsTensorModelDir(dir string) bool {
-	_, err := os.Stat(filepath.Join(dir, "model_index.json"))
-	return err == nil
-}
-
-// IsSafetensorsModelDir checks if the directory contains a standard safetensors model
-// by looking for config.json and at least one .safetensors file.
-func IsSafetensorsModelDir(dir string) bool {
-	// Must have config.json
-	if _, err := os.Stat(filepath.Join(dir, "config.json")); err != nil {
-		return false
-	}
-
-	// Must have at least one .safetensors file
-	entries, err := os.ReadDir(dir)
-	if err != nil {
-		return false
-	}
-
-	for _, entry := range entries {
-		if strings.HasSuffix(entry.Name(), ".safetensors") {
-			return true
-		}
-	}
-
-	return false
-}
-
-// LayerInfo holds metadata for a created layer.
-type LayerInfo struct {
-	Digest    string
-	Size      int64
-	MediaType string
-	Name      string // Path-style name: "component/tensor" or "path/to/config.json"
-}
-
-// LayerCreator is called to create a blob layer.
-// name is the path-style name (e.g., "tokenizer/tokenizer.json")
-type LayerCreator func(r io.Reader, mediaType, name string) (LayerInfo, error)
-
-// TensorLayerCreator creates a tensor blob layer with metadata.
-// name is the path-style name including component (e.g., "text_encoder/model.embed_tokens.weight")
-type TensorLayerCreator func(r io.Reader, name, dtype string, shape []int32) (LayerInfo, error)
-
-// QuantizingTensorLayerCreator creates tensor layers with optional quantization.
-// When quantize is non-empty (e.g., "fp8"), returns multiple layers (weight + scales + biases).
-type QuantizingTensorLayerCreator func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error)
-
-// ManifestWriter writes the manifest file.
-type ManifestWriter func(modelName string, config LayerInfo, layers []LayerInfo) error
-
-// ShouldQuantize returns true if a tensor should be quantized.
-// For image gen models (component non-empty): quantizes linear weights, skipping VAE, embeddings, norms.
-// For LLM models (component empty): quantizes linear weights, skipping embeddings, norms, and small tensors.
-func ShouldQuantize(name, component string) bool {
-	// Image gen specific: skip VAE entirely
-	if component == "vae" {
-		return false
-	}
-
-	// Skip embeddings
-	if strings.Contains(name, "embed") {
-		return false
-	}
-
-	// Skip layer norms and RMS norms
-	if strings.Contains(name, "norm") || strings.Contains(name, "ln_") || strings.Contains(name, "layernorm") {
-		return false
-	}
-
-	// Skip biases
-	if strings.HasSuffix(name, ".bias") {
-		return false
-	}
-
-	// Only quantize weights
-	return strings.HasSuffix(name, ".weight")
-}
-
-// ShouldQuantizeTensor returns true if a tensor should be quantized based on name and shape.
-// This is a more detailed check that also considers tensor dimensions.
-func ShouldQuantizeTensor(name string, shape []int32) bool {
-	// Use basic name-based check first
-	if !ShouldQuantize(name, "") {
-		return false
-	}
-
-	// Only quantize 2D tensors (linear layers) - skip 1D (biases, norms) and higher-D (convolutions if any)
-	if len(shape) != 2 {
-		return false
-	}
-
-	// Skip small tensors (less than 1024 elements) - not worth quantizing
-	if len(shape) >= 2 && int64(shape[0])*int64(shape[1]) < 1024 {
-		return false
-	}
-
-	// MLX quantization requires last dimension to be divisible by group size (32)
-	if shape[len(shape)-1]%32 != 0 {
-		return false
-	}
-
-	return true
-}
-
-// CreateSafetensorsModel imports a standard safetensors model from a directory.
-// This handles Hugging Face style models with config.json and *.safetensors files.
-// Stores each tensor as a separate blob for fine-grained deduplication.
-// If quantize is non-empty (e.g., "fp8"), eligible tensors will be quantized.
-func CreateSafetensorsModel(modelName, modelDir, quantize string, createLayer LayerCreator, createTensorLayer QuantizingTensorLayerCreator, writeManifest ManifestWriter, fn func(status string)) error {
-	var layers []LayerInfo
-	var configLayer LayerInfo
-
-	entries, err := os.ReadDir(modelDir)
-	if err != nil {
-		return fmt.Errorf("failed to read directory: %w", err)
-	}
-
-	// Process all safetensors files
-	for _, entry := range entries {
-		if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".safetensors") {
-			continue
-		}
-
-		stPath := filepath.Join(modelDir, entry.Name())
-
-		// Extract individual tensors from safetensors file
-		extractor, err := safetensors.OpenForExtraction(stPath)
-		if err != nil {
-			return fmt.Errorf("failed to open %s: %w", stPath, err)
-		}
-
-		tensorNames := extractor.ListTensors()
-		quantizeMsg := ""
-		if quantize != "" {
-			quantizeMsg = fmt.Sprintf(", quantizing to %s", quantize)
-		}
-		fn(fmt.Sprintf("importing %s (%d tensors%s)", entry.Name(), len(tensorNames), quantizeMsg))
-
-		for _, tensorName := range tensorNames {
-			td, err := extractor.GetTensor(tensorName)
-			if err != nil {
-				extractor.Close()
-				return fmt.Errorf("failed to get tensor %s: %w", tensorName, err)
-			}
-
-			// Determine quantization type for this tensor (empty string if not quantizing)
-			quantizeType := ""
-			if quantize != "" && ShouldQuantizeTensor(tensorName, td.Shape) {
-				quantizeType = quantize
-			}
-
-			// Store as minimal safetensors format (88 bytes header overhead)
-			// This enables native mmap loading via mlx_load_safetensors
-			// createTensorLayer returns multiple layers if quantizing (weight + scales)
-			newLayers, err := createTensorLayer(td.SafetensorsReader(), tensorName, td.Dtype, td.Shape, quantizeType)
-			if err != nil {
-				extractor.Close()
-				return fmt.Errorf("failed to create layer for %s: %w", tensorName, err)
-			}
-			layers = append(layers, newLayers...)
-		}
-
-		extractor.Close()
-	}
-
-	// Process all JSON config files
-	for _, entry := range entries {
-		if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".json") {
-			continue
-		}
-
-		// Skip the index file as we don't need it after extraction
-		if entry.Name() == "model.safetensors.index.json" {
-			continue
-		}
-
-		cfgPath := entry.Name()
-		fullPath := filepath.Join(modelDir, cfgPath)
-
-		fn(fmt.Sprintf("importing config %s", cfgPath))
-
-		f, err := os.Open(fullPath)
-		if err != nil {
-			return fmt.Errorf("failed to open %s: %w", cfgPath, err)
-		}
-
-		layer, err := createLayer(f, "application/vnd.ollama.image.json", cfgPath)
-		f.Close()
-		if err != nil {
-			return fmt.Errorf("failed to create layer for %s: %w", cfgPath, err)
-		}
-
-		// Use config.json as the config layer
-		if cfgPath == "config.json" {
-			configLayer = layer
-		}
-
-		layers = append(layers, layer)
-	}
-
-	if configLayer.Digest == "" {
-		return fmt.Errorf("config.json not found in %s", modelDir)
-	}
-
-	fn(fmt.Sprintf("writing manifest for %s", modelName))
-
-	if err := writeManifest(modelName, configLayer, layers); err != nil {
-		return fmt.Errorf("failed to write manifest: %w", err)
-	}
-
-	fn(fmt.Sprintf("successfully imported %s with %d layers", modelName, len(layers)))
-	return nil
-}

x/create/create_test.go

@@ -1,752 +0,0 @@
-package create
-
-import (
-	"bytes"
-	"encoding/binary"
-	"encoding/json"
-	"io"
-	"os"
-	"path/filepath"
-	"strings"
-	"testing"
-)
-
-func TestIsTensorModelDir(t *testing.T) {
-	tests := []struct {
-		name     string
-		setup    func(dir string) error
-		expected bool
-	}{
-		{
-			name: "valid diffusers model with model_index.json",
-			setup: func(dir string) error {
-				return os.WriteFile(filepath.Join(dir, "model_index.json"), []byte(`{"_class_name": "FluxPipeline"}`), 0o644)
-			},
-			expected: true,
-		},
-		{
-			name: "empty directory",
-			setup: func(dir string) error {
-				return nil
-			},
-			expected: false,
-		},
-		{
-			name: "directory with other files but no model_index.json",
-			setup: func(dir string) error {
-				return os.WriteFile(filepath.Join(dir, "config.json"), []byte(`{}`), 0o644)
-			},
-			expected: false,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			dir := t.TempDir()
-			if err := tt.setup(dir); err != nil {
-				t.Fatalf("setup failed: %v", err)
-			}
-
-			got := IsTensorModelDir(dir)
-			if got != tt.expected {
-				t.Errorf("IsTensorModelDir() = %v, want %v", got, tt.expected)
-			}
-		})
-	}
-}
-
-func TestIsSafetensorsModelDir(t *testing.T) {
-	tests := []struct {
-		name     string
-		setup    func(dir string) error
-		expected bool
-	}{
-		{
-			name: "valid safetensors model with config.json and .safetensors file",
-			setup: func(dir string) error {
-				if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(`{"model_type": "gemma3"}`), 0o644); err != nil {
-					return err
-				}
-				return os.WriteFile(filepath.Join(dir, "model.safetensors"), []byte("dummy"), 0o644)
-			},
-			expected: true,
-		},
-		{
-			name: "config.json only, no safetensors files",
-			setup: func(dir string) error {
-				return os.WriteFile(filepath.Join(dir, "config.json"), []byte(`{}`), 0o644)
-			},
-			expected: false,
-		},
-		{
-			name: "safetensors file only, no config.json",
-			setup: func(dir string) error {
-				return os.WriteFile(filepath.Join(dir, "model.safetensors"), []byte("dummy"), 0o644)
-			},
-			expected: false,
-		},
-		{
-			name: "empty directory",
-			setup: func(dir string) error {
-				return nil
-			},
-			expected: false,
-		},
-		{
-			name: "multiple safetensors files with config.json",
-			setup: func(dir string) error {
-				if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(`{}`), 0o644); err != nil {
-					return err
-				}
-				if err := os.WriteFile(filepath.Join(dir, "model-00001-of-00002.safetensors"), []byte("dummy"), 0o644); err != nil {
-					return err
-				}
-				return os.WriteFile(filepath.Join(dir, "model-00002-of-00002.safetensors"), []byte("dummy"), 0o644)
-			},
-			expected: true,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			dir := t.TempDir()
-			if err := tt.setup(dir); err != nil {
-				t.Fatalf("setup failed: %v", err)
-			}
-
-			got := IsSafetensorsModelDir(dir)
-			if got != tt.expected {
-				t.Errorf("IsSafetensorsModelDir() = %v, want %v", got, tt.expected)
-			}
-		})
-	}
-}
-
-func TestIsSafetensorsModelDir_NonexistentDir(t *testing.T) {
-	got := IsSafetensorsModelDir("/nonexistent/path/that/does/not/exist")
-	if got != false {
-		t.Errorf("IsSafetensorsModelDir() = %v for nonexistent dir, want false", got)
-	}
-}
-
-// createMinimalSafetensors creates a minimal valid safetensors file with one tensor
-func createMinimalSafetensors(t *testing.T, path string) {
-	t.Helper()
-
-	// Create a minimal safetensors file with a single float32 tensor
-	header := map[string]interface{}{
-		"test_tensor": map[string]interface{}{
-			"dtype":        "F32",
-			"shape":        []int{2, 2},
-			"data_offsets": []int{0, 16}, // 4 float32 values = 16 bytes
-		},
-	}
-	headerJSON, err := json.Marshal(header)
-	if err != nil {
-		t.Fatalf("failed to marshal header: %v", err)
-	}
-
-	// Pad header to 8-byte alignment
-	padding := (8 - len(headerJSON)%8) % 8
-	headerJSON = append(headerJSON, bytes.Repeat([]byte(" "), padding)...)
-
-	// Write file
-	f, err := os.Create(path)
-	if err != nil {
-		t.Fatalf("failed to create file: %v", err)
-	}
-	defer f.Close()
-
-	// Write header size (8 bytes, little endian)
-	if err := binary.Write(f, binary.LittleEndian, uint64(len(headerJSON))); err != nil {
-		t.Fatalf("failed to write header size: %v", err)
-	}
-
-	// Write header
-	if _, err := f.Write(headerJSON); err != nil {
-		t.Fatalf("failed to write header: %v", err)
-	}
-
-	// Write tensor data (16 bytes of zeros for 4 float32 values)
-	if _, err := f.Write(make([]byte, 16)); err != nil {
-		t.Fatalf("failed to write tensor data: %v", err)
-	}
-}
-
-func TestCreateSafetensorsModel(t *testing.T) {
-	dir := t.TempDir()
-
-	// Create config.json
-	configJSON := `{"model_type": "test", "architectures": ["TestModel"]}`
-	if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(configJSON), 0o644); err != nil {
-		t.Fatalf("failed to write config.json: %v", err)
-	}
-
-	// Create a minimal safetensors file
-	createMinimalSafetensors(t, filepath.Join(dir, "model.safetensors"))
-
-	// Track what was created
-	var createdLayers []LayerInfo
-	var manifestWritten bool
-	var manifestModelName string
-	var manifestConfigLayer LayerInfo
-	var manifestLayers []LayerInfo
-	var statusMessages []string
-
-	// Mock callbacks
-	createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) {
-		data, err := io.ReadAll(r)
-		if err != nil {
-			return LayerInfo{}, err
-		}
-		layer := LayerInfo{
-			Digest:    "sha256:test",
-			Size:      int64(len(data)),
-			MediaType: mediaType,
-			Name:      name,
-		}
-		createdLayers = append(createdLayers, layer)
-		return layer, nil
-	}
-
-	createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) {
-		data, err := io.ReadAll(r)
-		if err != nil {
-			return nil, err
-		}
-		layer := LayerInfo{
-			Digest:    "sha256:tensor_" + name,
-			Size:      int64(len(data)),
-			MediaType: "application/vnd.ollama.image.tensor",
-			Name:      name,
-		}
-		createdLayers = append(createdLayers, layer)
-		return []LayerInfo{layer}, nil
-	}
-
-	writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error {
-		manifestWritten = true
-		manifestModelName = modelName
-		manifestConfigLayer = config
-		manifestLayers = layers
-		return nil
-	}
-
-	progressFn := func(status string) {
-		statusMessages = append(statusMessages, status)
-	}
-
-	// Run CreateSafetensorsModel
-	err := CreateSafetensorsModel("test-model", dir, "", createLayer, createTensorLayer, writeManifest, progressFn)
-	if err != nil {
-		t.Fatalf("CreateSafetensorsModel failed: %v", err)
-	}
-
-	// Verify manifest was written
-	if !manifestWritten {
-		t.Error("manifest was not written")
-	}
-
-	if manifestModelName != "test-model" {
-		t.Errorf("manifest model name = %q, want %q", manifestModelName, "test-model")
-	}
-
-	// Verify config layer was set
-	if manifestConfigLayer.Name != "config.json" {
-		t.Errorf("config layer name = %q, want %q", manifestConfigLayer.Name, "config.json")
-	}
-
-	// Verify we have at least one tensor and one config layer
-	hasTensor := false
-	hasConfig := false
-	for _, layer := range manifestLayers {
-		if layer.Name == "test_tensor" {
-			hasTensor = true
-		}
-		if layer.Name == "config.json" {
-			hasConfig = true
-		}
-	}
-
-	if !hasTensor {
-		t.Error("no tensor layer found in manifest")
-	}
-	if !hasConfig {
-		t.Error("no config layer found in manifest")
-	}
-
-	// Verify status messages were sent
-	if len(statusMessages) == 0 {
-		t.Error("no status messages received")
-	}
-}
-
-func TestCreateSafetensorsModel_NoConfigJson(t *testing.T) {
-	dir := t.TempDir()
-
-	// Create only a safetensors file, no config.json
-	createMinimalSafetensors(t, filepath.Join(dir, "model.safetensors"))
-
-	// Mock callbacks (minimal)
-	createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) {
-		io.ReadAll(r)
-		return LayerInfo{Name: name}, nil
-	}
-	createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) {
-		io.ReadAll(r)
-		return []LayerInfo{{Name: name}}, nil
-	}
-	writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error {
-		return nil
-	}
-	progressFn := func(status string) {}
-
-	err := CreateSafetensorsModel("test-model", dir, "", createLayer, createTensorLayer, writeManifest, progressFn)
-	if err == nil {
-		t.Error("expected error for missing config.json, got nil")
-	}
-}
-
-func TestCreateSafetensorsModel_EmptyDir(t *testing.T) {
-	dir := t.TempDir()
-
-	// Mock callbacks
-	createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) {
-		return LayerInfo{}, nil
-	}
-	createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) {
-		return []LayerInfo{{}}, nil
-	}
-	writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error {
-		return nil
-	}
-	progressFn := func(status string) {}
-
-	err := CreateSafetensorsModel("test-model", dir, "", createLayer, createTensorLayer, writeManifest, progressFn)
-	if err == nil {
-		t.Error("expected error for empty directory, got nil")
-	}
-}
-
-func TestCreateSafetensorsModel_SkipsIndexJson(t *testing.T) {
-	dir := t.TempDir()
-
-	// Create config.json
-	if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(`{}`), 0o644); err != nil {
-		t.Fatalf("failed to write config.json: %v", err)
-	}
-
-	// Create model.safetensors.index.json (should be skipped)
-	indexJSON := `{"metadata": {"total_size": 100}, "weight_map": {}}`
-	if err := os.WriteFile(filepath.Join(dir, "model.safetensors.index.json"), []byte(indexJSON), 0o644); err != nil {
-		t.Fatalf("failed to write index.json: %v", err)
-	}
-
-	// Create a minimal safetensors file
-	createMinimalSafetensors(t, filepath.Join(dir, "model.safetensors"))
-
-	var configNames []string
-
-	createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) {
-		io.ReadAll(r)
-		configNames = append(configNames, name)
-		return LayerInfo{Name: name, Digest: "sha256:test"}, nil
-	}
-	createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) {
-		io.ReadAll(r)
-		return []LayerInfo{{Name: name}}, nil
-	}
-	writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error {
-		return nil
-	}
-	progressFn := func(status string) {}
-
-	err := CreateSafetensorsModel("test-model", dir, "", createLayer, createTensorLayer, writeManifest, progressFn)
-	if err != nil {
-		t.Fatalf("CreateSafetensorsModel failed: %v", err)
-	}
-
-	// Verify model.safetensors.index.json was not included
-	for _, name := range configNames {
-		if name == "model.safetensors.index.json" {
-			t.Error("model.safetensors.index.json should have been skipped")
-		}
-	}
-}
-
-func TestResolveManifestPath(t *testing.T) {
-	tests := []struct {
-		name      string
-		modelName string
-		wantParts []string // Parts that should appear in the path
-	}{
-		{
-			name:      "simple model name",
-			modelName: "llama2",
-			wantParts: []string{"registry.ollama.ai", "library", "llama2", "latest"},
-		},
-		{
-			name:      "model name with tag",
-			modelName: "llama2:7b",
-			wantParts: []string{"registry.ollama.ai", "library", "llama2", "7b"},
-		},
-		{
-			name:      "model name with namespace",
-			modelName: "myuser/mymodel",
-			wantParts: []string{"registry.ollama.ai", "myuser", "mymodel", "latest"},
-		},
-		{
-			name:      "model name with namespace and tag",
-			modelName: "myuser/mymodel:v1",
-			wantParts: []string{"registry.ollama.ai", "myuser", "mymodel", "v1"},
-		},
-		{
-			name:      "fully qualified model name",
-			modelName: "registry.example.com/namespace/model:tag",
-			wantParts: []string{"registry.example.com", "namespace", "model", "tag"},
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			got := resolveManifestPath(tt.modelName)
-
-			for _, part := range tt.wantParts {
-				if !strings.Contains(got, part) {
-					t.Errorf("resolveManifestPath(%q) = %q, missing part %q", tt.modelName, got, part)
-				}
-			}
-		})
-	}
-}
-
-func TestLayerInfo(t *testing.T) {
-	layer := LayerInfo{
-		Digest:    "sha256:abc123",
-		Size:      1024,
-		MediaType: "application/vnd.ollama.image.tensor",
-		Name:      "model.weight",
-	}
-
-	if layer.Digest != "sha256:abc123" {
-		t.Errorf("Digest = %q, want %q", layer.Digest, "sha256:abc123")
-	}
-	if layer.Size != 1024 {
-		t.Errorf("Size = %d, want %d", layer.Size, 1024)
-	}
-	if layer.MediaType != "application/vnd.ollama.image.tensor" {
-		t.Errorf("MediaType = %q, want %q", layer.MediaType, "application/vnd.ollama.image.tensor")
-	}
-	if layer.Name != "model.weight" {
-		t.Errorf("Name = %q, want %q", layer.Name, "model.weight")
-	}
-}
-
-func TestModelConfig(t *testing.T) {
-	config := ModelConfig{
-		ModelFormat:  "safetensors",
-		Capabilities: []string{"completion", "chat"},
-	}
-
-	if config.ModelFormat != "safetensors" {
-		t.Errorf("ModelFormat = %q, want %q", config.ModelFormat, "safetensors")
-	}
-	if len(config.Capabilities) != 2 {
-		t.Errorf("Capabilities length = %d, want %d", len(config.Capabilities), 2)
-	}
-}
-
-func TestManifest(t *testing.T) {
-	manifest := Manifest{
-		SchemaVersion: 2,
-		MediaType:     "application/vnd.oci.image.manifest.v1+json",
-		Config: ManifestLayer{
-			MediaType: "application/vnd.docker.container.image.v1+json",
-			Digest:    "sha256:config",
-			Size:      100,
-		},
-		Layers: []ManifestLayer{
-			{
-				MediaType: "application/vnd.ollama.image.tensor",
-				Digest:    "sha256:layer1",
-				Size:      1000,
-				Name:      "weight.bin",
-			},
-		},
-	}
-
-	if manifest.SchemaVersion != 2 {
-		t.Errorf("SchemaVersion = %d, want %d", manifest.SchemaVersion, 2)
-	}
-	if manifest.Config.Digest != "sha256:config" {
-		t.Errorf("Config.Digest = %q, want %q", manifest.Config.Digest, "sha256:config")
-	}
-	if len(manifest.Layers) != 1 {
-		t.Errorf("Layers length = %d, want %d", len(manifest.Layers), 1)
-	}
-	if manifest.Layers[0].Name != "weight.bin" {
-		t.Errorf("Layers[0].Name = %q, want %q", manifest.Layers[0].Name, "weight.bin")
-	}
-}
-
-func TestShouldQuantize(t *testing.T) {
-	tests := []struct {
-		name      string
-		tensor    string
-		component string
-		want      bool
-	}{
-		// VAE component should never be quantized
-		{"vae weight", "decoder.weight", "vae", false},
-		{"vae bias", "decoder.bias", "vae", false},
-
-		// Embeddings should not be quantized
-		{"embedding weight", "embed_tokens.weight", "", false},
-		{"embedding in name", "token_embedding.weight", "", false},
-
-		// Norms should not be quantized
-		{"layer norm", "layer_norm.weight", "", false},
-		{"rms norm", "rms_norm.weight", "", false},
-		{"ln prefix", "ln_1.weight", "", false},
-		{"layernorm in name", "input_layernorm.weight", "", false},
-
-		// Biases should not be quantized
-		{"bias tensor", "attention.bias", "", false},
-		{"proj bias", "o_proj.bias", "", false},
-
-		// Linear weights should be quantized
-		{"linear weight", "q_proj.weight", "", true},
-		{"attention weight", "self_attn.weight", "", true},
-		{"mlp weight", "mlp.gate_proj.weight", "", true},
-
-		// Transformer component weights should be quantized
-		{"transformer weight", "layers.0.weight", "transformer", true},
-		{"text_encoder weight", "encoder.weight", "text_encoder", true},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			got := ShouldQuantize(tt.tensor, tt.component)
-			if got != tt.want {
-				t.Errorf("ShouldQuantize(%q, %q) = %v, want %v", tt.tensor, tt.component, got, tt.want)
-			}
-		})
-	}
-}
-
-func TestShouldQuantizeTensor(t *testing.T) {
-	tests := []struct {
-		name   string
-		tensor string
-		shape  []int32
-		want   bool
-	}{
-		// 2D tensors with sufficient size should be quantized
-		{"large 2D weight", "q_proj.weight", []int32{4096, 4096}, true},
-		{"medium 2D weight", "small_proj.weight", []int32{128, 128}, true},
-
-		// Small tensors should not be quantized (< 1024 elements)
-		{"tiny 2D weight", "tiny.weight", []int32{16, 16}, false},
-		{"small 2D weight", "small.weight", []int32{31, 31}, false},
-
-		// 1D tensors should not be quantized
-		{"1D tensor", "layer_norm.weight", []int32{4096}, false},
-
-		// 3D+ tensors should not be quantized
-		{"3D tensor", "conv.weight", []int32{64, 64, 3}, false},
-		{"4D tensor", "conv2d.weight", []int32{64, 64, 3, 3}, false},
-
-		// Embeddings should not be quantized regardless of shape
-		{"embedding 2D", "embed_tokens.weight", []int32{32000, 4096}, false},
-
-		// Norms should not be quantized regardless of shape
-		{"norm 2D", "layer_norm.weight", []int32{4096, 1}, false},
-
-		// Biases should not be quantized
-		{"bias 2D", "proj.bias", []int32{4096, 1}, false},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			got := ShouldQuantizeTensor(tt.tensor, tt.shape)
-			if got != tt.want {
-				t.Errorf("ShouldQuantizeTensor(%q, %v) = %v, want %v", tt.tensor, tt.shape, got, tt.want)
-			}
-		})
-	}
-}
-
-func TestCreateSafetensorsModel_WithQuantize(t *testing.T) {
-	dir := t.TempDir()
-
-	// Create config.json
-	configJSON := `{"model_type": "test", "architectures": ["TestModel"]}`
-	if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(configJSON), 0o644); err != nil {
-		t.Fatalf("failed to write config.json: %v", err)
-	}
-
-	// Create a minimal safetensors file
-	createMinimalSafetensors(t, filepath.Join(dir, "model.safetensors"))
-
-	var quantizeRequested []string
-
-	createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) {
-		io.ReadAll(r)
-		return LayerInfo{Name: name, Digest: "sha256:test"}, nil
-	}
-
-	createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) {
-		io.ReadAll(r)
-		quantizeRequested = append(quantizeRequested, quantize)
-		return []LayerInfo{{Name: name}}, nil
-	}
-
-	writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error {
-		return nil
-	}
-
-	progressFn := func(status string) {}
-
-	// Run with quantize enabled
-	err := CreateSafetensorsModel("test-model", dir, "fp8", createLayer, createTensorLayer, writeManifest, progressFn)
-	if err != nil {
-		t.Fatalf("CreateSafetensorsModel failed: %v", err)
-	}
-
-	// Verify quantize was passed to callback (will be false for small test tensor)
-	if len(quantizeRequested) == 0 {
-		t.Error("no tensors processed")
-	}
-}
-
-// createMinimalImageGenModel creates a minimal diffusers-style model directory
-func createMinimalImageGenModel(t *testing.T, dir string) {
-	t.Helper()
-
-	// Create model_index.json
-	modelIndex := `{"_class_name": "FluxPipeline", "_diffusers_version": "0.30.0"}`
-	if err := os.WriteFile(filepath.Join(dir, "model_index.json"), []byte(modelIndex), 0o644); err != nil {
-		t.Fatalf("failed to write model_index.json: %v", err)
-	}
-
-	// Create transformer directory with a safetensors file
-	transformerDir := filepath.Join(dir, "transformer")
-	if err := os.MkdirAll(transformerDir, 0o755); err != nil {
-		t.Fatalf("failed to create transformer dir: %v", err)
-	}
-	createMinimalSafetensors(t, filepath.Join(transformerDir, "model.safetensors"))
-
-	// Create transformer config
-	transformerConfig := `{"hidden_size": 3072}`
-	if err := os.WriteFile(filepath.Join(transformerDir, "config.json"), []byte(transformerConfig), 0o644); err != nil {
-		t.Fatalf("failed to write transformer config: %v", err)
-	}
-}
-
-func TestCreateImageGenModel(t *testing.T) {
-	dir := t.TempDir()
-	createMinimalImageGenModel(t, dir)
-
-	var manifestWritten bool
-	var manifestModelName string
-	var statusMessages []string
-
-	createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) {
-		io.ReadAll(r)
-		return LayerInfo{Name: name, Digest: "sha256:test"}, nil
-	}
-
-	createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) {
-		io.ReadAll(r)
-		return []LayerInfo{{Name: name, Digest: "sha256:tensor"}}, nil
-	}
-
-	writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error {
-		manifestWritten = true
-		manifestModelName = modelName
-		return nil
-	}
-
-	progressFn := func(status string) {
-		statusMessages = append(statusMessages, status)
-	}
-
-	err := CreateImageGenModel("test-imagegen", dir, "", createLayer, createTensorLayer, writeManifest, progressFn)
-	if err != nil {
-		t.Fatalf("CreateImageGenModel failed: %v", err)
-	}
-
-	if !manifestWritten {
-		t.Error("manifest was not written")
-	}
-
-	if manifestModelName != "test-imagegen" {
-		t.Errorf("manifest model name = %q, want %q", manifestModelName, "test-imagegen")
-	}
-
-	if len(statusMessages) == 0 {
-		t.Error("no status messages received")
-	}
-}
-
-func TestCreateImageGenModel_NoModelIndex(t *testing.T) {
-	dir := t.TempDir()
-
-	// Create only transformer without model_index.json
-	transformerDir := filepath.Join(dir, "transformer")
-	if err := os.MkdirAll(transformerDir, 0o755); err != nil {
-		t.Fatalf("failed to create transformer dir: %v", err)
-	}
-	createMinimalSafetensors(t, filepath.Join(transformerDir, "model.safetensors"))
-
-	createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) {
-		io.ReadAll(r)
-		return LayerInfo{Name: name}, nil
-	}
-	createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) {
-		io.ReadAll(r)
-		return []LayerInfo{{Name: name}}, nil
-	}
-	writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error {
-		return nil
-	}
-	progressFn := func(status string) {}
-
-	err := CreateImageGenModel("test-imagegen", dir, "", createLayer, createTensorLayer, writeManifest, progressFn)
-	if err == nil {
-		t.Error("expected error for missing model_index.json, got nil")
-	}
-}
-
-func TestCreateImageGenModel_WithQuantize(t *testing.T) {
-	dir := t.TempDir()
-	createMinimalImageGenModel(t, dir)
-
-	var quantizeRequested []string
-
-	createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) {
-		io.ReadAll(r)
-		return LayerInfo{Name: name, Digest: "sha256:test"}, nil
-	}
-
-	createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) {
-		io.ReadAll(r)
-		quantizeRequested = append(quantizeRequested, quantize)
-		return []LayerInfo{{Name: name}}, nil
-	}
-
-	writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error {
-		return nil
-	}
-
-	progressFn := func(status string) {}
-
-	err := CreateImageGenModel("test-imagegen", dir, "fp8", createLayer, createTensorLayer, writeManifest, progressFn)
-	if err != nil {
-		t.Fatalf("CreateImageGenModel failed: %v", err)
-	}
-
-	if len(quantizeRequested) == 0 {
-		t.Error("no tensors processed")
-	}
-}

x/imagegen/api/handler.go

@@ -0,0 +1,231 @@
+package api
+
+import (
+	"fmt"
+	"net/http"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/gin-gonic/gin"
+
+	"github.com/ollama/ollama/api"
+	"github.com/ollama/ollama/llm"
+	"github.com/ollama/ollama/x/imagegen"
+)
+
+// RunnerScheduler is the interface for scheduling a model runner.
+// This is implemented by server.Server to avoid circular imports.
+type RunnerScheduler interface {
+	ScheduleImageGenRunner(ctx *gin.Context, modelName string, opts api.Options, keepAlive *api.Duration) (llm.LlamaServer, error)
+}
+
+// RegisterRoutes registers the image generation API routes.
+func RegisterRoutes(r gin.IRouter, scheduler RunnerScheduler) {
+	r.POST("/v1/images/generations", func(c *gin.Context) {
+		ImageGenerationHandler(c, scheduler)
+	})
+}
+
+// ImageGenerationHandler handles OpenAI-compatible image generation requests.
+func ImageGenerationHandler(c *gin.Context, scheduler RunnerScheduler) {
+	var req ImageGenerationRequest
+	if err := c.BindJSON(&req); err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{"message": err.Error()}})
+		return
+	}
+
+	// Validate required fields
+	if req.Model == "" {
+		c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{"message": "model is required"}})
+		return
+	}
+	if req.Prompt == "" {
+		c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{"message": "prompt is required"}})
+		return
+	}
+
+	// Apply defaults
+	if req.N == 0 {
+		req.N = 1
+	}
+	if req.Size == "" {
+		req.Size = "1024x1024"
+	}
+	if req.ResponseFormat == "" {
+		req.ResponseFormat = "b64_json"
+	}
+
+	// Verify model exists
+	if imagegen.ResolveModelName(req.Model) == "" {
+		c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"message": fmt.Sprintf("model %q not found", req.Model)}})
+		return
+	}
+
+	// Parse size
+	width, height := parseSize(req.Size)
+
+	// Build options - we repurpose NumCtx/NumGPU for width/height
+	opts := api.Options{}
+	opts.NumCtx = int(width)
+	opts.NumGPU = int(height)
+
+	// Schedule runner
+	runner, err := scheduler.ScheduleImageGenRunner(c, req.Model, opts, nil)
+	if err != nil {
+		status := http.StatusInternalServerError
+		if strings.Contains(err.Error(), "not found") {
+			status = http.StatusNotFound
+		}
+		c.JSON(status, gin.H{"error": gin.H{"message": err.Error()}})
+		return
+	}
+
+	// Build completion request
+	completionReq := llm.CompletionRequest{
+		Prompt:  req.Prompt,
+		Options: &opts,
+	}
+
+	if req.Stream {
+		handleStreamingResponse(c, runner, completionReq, req.ResponseFormat)
+	} else {
+		handleNonStreamingResponse(c, runner, completionReq, req.ResponseFormat)
+	}
+}
+
+func handleStreamingResponse(c *gin.Context, runner llm.LlamaServer, req llm.CompletionRequest, format string) {
+	c.Header("Content-Type", "text/event-stream")
+	c.Header("Cache-Control", "no-cache")
+	c.Header("Connection", "keep-alive")
+
+	var imageBase64 string
+	err := runner.Completion(c.Request.Context(), req, func(resp llm.CompletionResponse) {
+		if resp.Done {
+			imageBase64 = extractBase64(resp.Content)
+		} else {
+			progress := parseProgress(resp.Content)
+			if progress.Total > 0 {
+				c.SSEvent("progress", progress)
+				c.Writer.Flush()
+			}
+		}
+	})
+	if err != nil {
+		c.SSEvent("error", gin.H{"error": err.Error()})
+		return
+	}
+
+	c.SSEvent("done", buildResponse(imageBase64, format))
+}
+
+func handleNonStreamingResponse(c *gin.Context, runner llm.LlamaServer, req llm.CompletionRequest, format string) {
+	var imageBase64 string
+	err := runner.Completion(c.Request.Context(), req, func(resp llm.CompletionResponse) {
+		if resp.Done {
+			imageBase64 = extractBase64(resp.Content)
+		}
+	})
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"error": gin.H{"message": err.Error()}})
+		return
+	}
+
+	c.JSON(http.StatusOK, buildResponse(imageBase64, format))
+}
+
+func parseSize(size string) (int32, int32) {
+	parts := strings.Split(size, "x")
+	if len(parts) != 2 {
+		return 1024, 1024
+	}
+	w, _ := strconv.Atoi(parts[0])
+	h, _ := strconv.Atoi(parts[1])
+	if w == 0 {
+		w = 1024
+	}
+	if h == 0 {
+		h = 1024
+	}
+	return int32(w), int32(h)
+}
+
+func extractBase64(content string) string {
+	if strings.HasPrefix(content, "IMAGE_BASE64:") {
+		return content[13:]
+	}
+	return ""
+}
+
+func parseProgress(content string) ImageProgressEvent {
+	var step, total int
+	fmt.Sscanf(content, "\rGenerating: step %d/%d", &step, &total)
+	return ImageProgressEvent{Step: step, Total: total}
+}
+
+func buildResponse(imageBase64, format string) ImageGenerationResponse {
+	resp := ImageGenerationResponse{
+		Created: time.Now().Unix(),
+		Data:    make([]ImageData, 1),
+	}
+
+	if imageBase64 == "" {
+		return resp
+	}
+
+	if format == "url" {
+		// URL format not supported when using base64 transfer
+		resp.Data[0].B64JSON = imageBase64
+	} else {
+		resp.Data[0].B64JSON = imageBase64
+	}
+
+	return resp
+}
+
+// HandleGenerateRequest handles Ollama /api/generate requests for image gen models.
+// This allows routes.go to delegate image generation with minimal code.
+func HandleGenerateRequest(c *gin.Context, scheduler RunnerScheduler, modelName, prompt string, keepAlive *api.Duration, streamFn func(c *gin.Context, ch chan any)) {
+	opts := api.Options{}
+
+	// Schedule runner
+	runner, err := scheduler.ScheduleImageGenRunner(c, modelName, opts, keepAlive)
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
+		return
+	}
+
+	// Build completion request
+	completionReq := llm.CompletionRequest{
+		Prompt:  prompt,
+		Options: &opts,
+	}
+
+	// Stream responses via channel
+	ch := make(chan any)
+	go func() {
+		defer close(ch)
+		err := runner.Completion(c.Request.Context(), completionReq, func(resp llm.CompletionResponse) {
+			ch <- GenerateResponse{
+				Model:     modelName,
+				CreatedAt: time.Now().UTC(),
+				Response:  resp.Content,
+				Done:      resp.Done,
+			}
+		})
+		if err != nil {
+			// Log error but don't block - channel is already being consumed
+			_ = err
+		}
+	}()
+
+	streamFn(c, ch)
+}
+
+// GenerateResponse matches api.GenerateResponse structure for streaming.
+type GenerateResponse struct {
+	Model     string    `json:"model"`
+	CreatedAt time.Time `json:"created_at"`
+	Response  string    `json:"response"`
+	Done      bool      `json:"done"`
+}

x/imagegen/api/types.go

@@ -0,0 +1,31 @@
+// Package api provides OpenAI-compatible image generation API types.
+package api
+
+// ImageGenerationRequest is an OpenAI-compatible image generation request.
+type ImageGenerationRequest struct {
+	Model          string `json:"model"`
+	Prompt         string `json:"prompt"`
+	N              int    `json:"n,omitempty"`
+	Size           string `json:"size,omitempty"`
+	ResponseFormat string `json:"response_format,omitempty"`
+	Stream         bool   `json:"stream,omitempty"`
+}
+
+// ImageGenerationResponse is an OpenAI-compatible image generation response.
+type ImageGenerationResponse struct {
+	Created int64       `json:"created"`
+	Data    []ImageData `json:"data"`
+}
+
+// ImageData contains the generated image data.
+type ImageData struct {
+	URL           string `json:"url,omitempty"`
+	B64JSON       string `json:"b64_json,omitempty"`
+	RevisedPrompt string `json:"revised_prompt,omitempty"`
+}
+
+// ImageProgressEvent is sent during streaming to indicate generation progress.
+type ImageProgressEvent struct {
+	Step  int `json:"step"`
+	Total int `json:"total"`
+}

x/imagegen/cli.go

@@ -7,6 +7,7 @@ package imagegen
 
 import (
 	"encoding/base64"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
@@ -38,20 +39,79 @@ func DefaultOptions() ImageGenOptions {
 	return ImageGenOptions{
 		Width:  1024,
 		Height: 1024,
-		Steps:  0, // 0 means model default
+		Steps:  9,
 		Seed:   0, // 0 means random
 	}
 }
 
+// ModelInfo contains metadata about an image generation model.
+type ModelInfo struct {
+	Architecture   string
+	ParameterCount int64
+	Quantization   string
+}
+
+// GetModelInfo returns metadata about an image generation model.
+func GetModelInfo(modelName string) (*ModelInfo, error) {
+	manifest, err := LoadManifest(modelName)
+	if err != nil {
+		return nil, fmt.Errorf("failed to load manifest: %w", err)
+	}
+
+	info := &ModelInfo{}
+
+	// Read model_index.json for architecture, parameter count, and quantization
+	if data, err := manifest.ReadConfig("model_index.json"); err == nil {
+		var index struct {
+			Architecture   string `json:"architecture"`
+			ParameterCount int64  `json:"parameter_count"`
+			Quantization   string `json:"quantization"`
+		}
+		if json.Unmarshal(data, &index) == nil {
+			info.Architecture = index.Architecture
+			info.ParameterCount = index.ParameterCount
+			info.Quantization = index.Quantization
+		}
+	}
+
+	// Fallback: detect quantization from tensor names if not in config
+	if info.Quantization == "" {
+		for _, layer := range manifest.Manifest.Layers {
+			if strings.HasSuffix(layer.Name, ".weight_scale") {
+				info.Quantization = "FP8"
+				break
+			}
+		}
+		if info.Quantization == "" {
+			info.Quantization = "BF16"
+		}
+	}
+
+	// Fallback: estimate parameter count if not in config
+	if info.ParameterCount == 0 {
+		var totalSize int64
+		for _, layer := range manifest.Manifest.Layers {
+			if layer.MediaType == "application/vnd.ollama.image.tensor" {
+				if !strings.HasSuffix(layer.Name, "_scale") && !strings.HasSuffix(layer.Name, "_qbias") {
+					totalSize += layer.Size
+				}
+			}
+		}
+		// Assume BF16 (2 bytes/param) as rough estimate
+		info.ParameterCount = totalSize / 2
+	}
+
+	return info, nil
+}
+
 // RegisterFlags adds image generation flags to the given command.
 // Flags are hidden since they only apply to image generation models.
 func RegisterFlags(cmd *cobra.Command) {
 	cmd.Flags().Int("width", 1024, "Image width")
 	cmd.Flags().Int("height", 1024, "Image height")
-	cmd.Flags().Int("steps", 0, "Denoising steps (0 = model default)")
+	cmd.Flags().Int("steps", 9, "Denoising steps")
 	cmd.Flags().Int("seed", 0, "Random seed (0 for random)")
 	cmd.Flags().String("negative", "", "Negative prompt")
-	// Hide from main flags section - shown in separate section via AppendFlagsDocs
 	cmd.Flags().MarkHidden("width")
 	cmd.Flags().MarkHidden("height")
 	cmd.Flags().MarkHidden("steps")
@@ -59,19 +119,6 @@ func RegisterFlags(cmd *cobra.Command) {
 	cmd.Flags().MarkHidden("negative")
 }
 
-// AppendFlagsDocs appends image generation flags documentation to the command's usage template.
-func AppendFlagsDocs(cmd *cobra.Command) {
-	usage := `
-Image Generation Flags (experimental):
-      --width int      Image width
-      --height int     Image height
-      --steps int      Denoising steps
-      --seed int       Random seed
-      --negative str   Negative prompt
-`
-	cmd.SetUsageTemplate(cmd.UsageTemplate() + usage)
-}
-
 // RunCLI handles the CLI for image generation models.
 // Returns true if it handled the request, false if the caller should continue with normal flow.
 // Supports flags: --width, --height, --steps, --seed, --negative
@@ -111,15 +158,17 @@ func generateImageWithOptions(cmd *cobra.Command, modelName, prompt string, keep
 		return err
 	}
 
+	// Build request with image gen options encoded in Options fields
+	// NumCtx=width, NumGPU=height, NumPredict=steps, Seed=seed
 	req := &api.GenerateRequest{
 		Model:  modelName,
 		Prompt: prompt,
-		Width:  int32(opts.Width),
-		Height: int32(opts.Height),
-		Steps:  int32(opts.Steps),
-	}
-	if opts.Seed != 0 {
-		req.Options = map[string]any{"seed": opts.Seed}
+		Options: map[string]any{
+			"num_ctx":     opts.Width,
+			"num_gpu":     opts.Height,
+			"num_predict": opts.Steps,
+			"seed":        opts.Seed,
+		},
 	}
 	if keepAlive != nil {
 		req.KeepAlive = keepAlive
@@ -133,25 +182,32 @@ func generateImageWithOptions(cmd *cobra.Command, modelName, prompt string, keep
 	var stepBar *progress.StepBar
 	var imageBase64 string
 	err = client.Generate(cmd.Context(), req, func(resp api.GenerateResponse) error {
-		// Handle progress updates using structured fields
-		if resp.Total > 0 {
-			if stepBar == nil {
+		content := resp.Response
+
+		// Handle progress updates - parse step info and switch to step bar
+		if strings.HasPrefix(content, "\rGenerating:") {
+			var step, total int
+			fmt.Sscanf(content, "\rGenerating: step %d/%d", &step, &total)
+			if stepBar == nil && total > 0 {
 				spinner.Stop()
-				stepBar = progress.NewStepBar("Generating", int(resp.Total))
+				stepBar = progress.NewStepBar("Generating", total)
 				p.Add("", stepBar)
 			}
-			stepBar.Set(int(resp.Completed))
+			if stepBar != nil {
+				stepBar.Set(step)
+			}
+			return nil
 		}
 
-		// Handle final response with image data
-		if resp.Done && resp.Image != "" {
-			imageBase64 = resp.Image
+		// Handle final response with base64 image data
+		if resp.Done && strings.HasPrefix(content, "IMAGE_BASE64:") {
+			imageBase64 = content[13:]
 		}
 
 		return nil
 	})
 
-	p.StopAndClear()
+	p.Stop()
 	if err != nil {
 		return err
 	}
@@ -189,23 +245,6 @@ func runInteractive(cmd *cobra.Command, modelName string, keepAlive *api.Duratio
 		return err
 	}
 
-	// Preload the model with the specified keepalive
-	p := progress.NewProgress(os.Stderr)
-	spinner := progress.NewSpinner("")
-	p.Add("", spinner)
-
-	preloadReq := &api.GenerateRequest{
-		Model:     modelName,
-		KeepAlive: keepAlive,
-	}
-	if err := client.Generate(cmd.Context(), preloadReq, func(resp api.GenerateResponse) error {
-		return nil
-	}); err != nil {
-		p.StopAndClear()
-		return fmt.Errorf("failed to load model: %w", err)
-	}
-	p.StopAndClear()
-
 	scanner, err := readline.New(readline.Prompt{
 		Prompt:      ">>> ",
 		Placeholder: "Describe an image to generate (/help for commands)",
@@ -243,7 +282,7 @@ func runInteractive(cmd *cobra.Command, modelName string, keepAlive *api.Duratio
 		case strings.HasPrefix(line, "/bye"):
 			return nil
 		case strings.HasPrefix(line, "/?"), strings.HasPrefix(line, "/help"):
-			printInteractiveHelp()
+			printInteractiveHelp(opts)
 			continue
 		case strings.HasPrefix(line, "/set "):
 			if err := handleSetCommand(line[5:], &opts); err != nil {
@@ -262,12 +301,12 @@ func runInteractive(cmd *cobra.Command, modelName string, keepAlive *api.Duratio
 		req := &api.GenerateRequest{
 			Model:  modelName,
 			Prompt: line,
-			Width:  int32(opts.Width),
-			Height: int32(opts.Height),
-			Steps:  int32(opts.Steps),
-		}
-		if opts.Seed != 0 {
-			req.Options = map[string]any{"seed": opts.Seed}
+			Options: map[string]any{
+				"num_ctx":     opts.Width,
+				"num_gpu":     opts.Height,
+				"num_predict": opts.Steps,
+				"seed":        opts.Seed,
+			},
 		}
 		if keepAlive != nil {
 			req.KeepAlive = keepAlive
@@ -282,25 +321,32 @@ func runInteractive(cmd *cobra.Command, modelName string, keepAlive *api.Duratio
 		var imageBase64 string
 
 		err = client.Generate(cmd.Context(), req, func(resp api.GenerateResponse) error {
-			// Handle progress updates using structured fields
-			if resp.Total > 0 {
-				if stepBar == nil {
+			content := resp.Response
+
+			// Handle progress updates - parse step info and switch to step bar
+			if strings.HasPrefix(content, "\rGenerating:") {
+				var step, total int
+				fmt.Sscanf(content, "\rGenerating: step %d/%d", &step, &total)
+				if stepBar == nil && total > 0 {
 					spinner.Stop()
-					stepBar = progress.NewStepBar("Generating", int(resp.Total))
+					stepBar = progress.NewStepBar("Generating", total)
 					p.Add("", stepBar)
 				}
-				stepBar.Set(int(resp.Completed))
+				if stepBar != nil {
+					stepBar.Set(step)
+				}
+				return nil
 			}
 
-			// Handle final response with image data
-			if resp.Done && resp.Image != "" {
-				imageBase64 = resp.Image
+			// Handle final response with base64 image data
+			if resp.Done && strings.HasPrefix(content, "IMAGE_BASE64:") {
+				imageBase64 = content[13:]
 			}
 
 			return nil
 		})
 
-		p.StopAndClear()
+		p.Stop()
 		if err != nil {
 			fmt.Fprintf(os.Stderr, "Error: %v\n", err)
 			continue
@@ -351,13 +397,12 @@ func sanitizeFilename(s string) string {
 }
 
 // printInteractiveHelp prints help for interactive mode commands.
-// TODO: reconcile /set commands with /set parameter in text gen REPL (cmd/cmd.go)
-func printInteractiveHelp() {
+func printInteractiveHelp(opts ImageGenOptions) {
 	fmt.Fprintln(os.Stderr, "Commands:")
-	fmt.Fprintln(os.Stderr, "  /set width <n>     Set image width")
-	fmt.Fprintln(os.Stderr, "  /set height <n>    Set image height")
-	fmt.Fprintln(os.Stderr, "  /set steps <n>     Set denoising steps")
-	fmt.Fprintln(os.Stderr, "  /set seed <n>      Set random seed")
+	fmt.Fprintln(os.Stderr, "  /set width <n>     Set image width (current:", opts.Width, ")")
+	fmt.Fprintln(os.Stderr, "  /set height <n>    Set image height (current:", opts.Height, ")")
+	fmt.Fprintln(os.Stderr, "  /set steps <n>     Set denoising steps (current:", opts.Steps, ")")
+	fmt.Fprintln(os.Stderr, "  /set seed <n>      Set random seed (current:", opts.Seed, ", 0=random)")
 	fmt.Fprintln(os.Stderr, "  /set negative <s>  Set negative prompt")
 	fmt.Fprintln(os.Stderr, "  /show              Show current settings")
 	fmt.Fprintln(os.Stderr, "  /bye               Exit")

x/imagegen/client/create.go

@@ -0,0 +1,190 @@
+// Package client provides client-side model creation for tensor-based models.
+//
+// This package is in x/ because the tensor model storage format is under development.
+// It also exists to break an import cycle: server imports x/imagegen, so x/imagegen
+// cannot import server. This sub-package can import server because server doesn't
+// import it.
+//
+// TODO (jmorganca): This is temporary. When tensor models are promoted to production:
+//  1. Add proper API endpoints for tensor model creation
+//  2. Move tensor extraction to server-side
+//  3. Remove this package
+//  4. Follow the same client→server pattern as regular model creation
+package client
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"io"
+
+	"github.com/ollama/ollama/progress"
+	"github.com/ollama/ollama/server"
+	"github.com/ollama/ollama/types/model"
+	"github.com/ollama/ollama/x/imagegen"
+)
+
+// MinOllamaVersion is the minimum Ollama version required for image generation models.
+const MinOllamaVersion = "0.14.0"
+
+// CreateModel imports a tensor-based model from a local directory.
+// This creates blobs and manifest directly on disk, bypassing the HTTP API.
+// If quantize is "fp8", weights will be quantized to mxfp8 format during import.
+//
+// TODO (jmorganca): Replace with API-based creation when promoted to production.
+func CreateModel(modelName, modelDir, quantize string, p *progress.Progress) error {
+	if !imagegen.IsTensorModelDir(modelDir) {
+		return fmt.Errorf("%s is not an image generation model directory (model_index.json not found)", modelDir)
+	}
+
+	status := "importing image generation model"
+	spinner := progress.NewSpinner(status)
+	p.Add("imagegen", spinner)
+
+	// Create layer callback for config files
+	createLayer := func(r io.Reader, mediaType, name string) (imagegen.LayerInfo, error) {
+		layer, err := server.NewLayer(r, mediaType)
+		if err != nil {
+			return imagegen.LayerInfo{}, err
+		}
+		layer.Name = name
+
+		return imagegen.LayerInfo{
+			Digest:    layer.Digest,
+			Size:      layer.Size,
+			MediaType: layer.MediaType,
+			Name:      name,
+		}, nil
+	}
+
+	// Create tensor layer callback for individual tensors
+	// name is path-style: "component/tensor_name"
+	// When quantize is true, returns multiple layers (weight + scales)
+	createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, doQuantize bool) ([]imagegen.LayerInfo, error) {
+		if doQuantize {
+			// Check if quantization is supported
+			if !QuantizeSupported() {
+				return nil, fmt.Errorf("quantization requires MLX support")
+			}
+
+			// Quantize the tensor (affine mode returns weight, scales, qbiases)
+			qweightData, scalesData, qbiasData, _, _, _, err := quantizeTensor(r, name, dtype, shape)
+			if err != nil {
+				return nil, fmt.Errorf("failed to quantize %s: %w", name, err)
+			}
+
+			// Create layer for quantized weight
+			weightLayer, err := server.NewLayer(bytes.NewReader(qweightData), server.MediaTypeImageTensor)
+			if err != nil {
+				return nil, err
+			}
+
+			// Create layer for scales (use _scale suffix convention)
+			scalesLayer, err := server.NewLayer(bytes.NewReader(scalesData), server.MediaTypeImageTensor)
+			if err != nil {
+				return nil, err
+			}
+
+			layers := []imagegen.LayerInfo{
+				{
+					Digest:    weightLayer.Digest,
+					Size:      weightLayer.Size,
+					MediaType: weightLayer.MediaType,
+					Name:      name, // Keep original name for weight
+				},
+				{
+					Digest:    scalesLayer.Digest,
+					Size:      scalesLayer.Size,
+					MediaType: scalesLayer.MediaType,
+					Name:      name + "_scale", // Add _scale suffix
+				},
+			}
+
+			// Add qbiases layer if present (affine mode)
+			if qbiasData != nil {
+				qbiasLayer, err := server.NewLayer(bytes.NewReader(qbiasData), server.MediaTypeImageTensor)
+				if err != nil {
+					return nil, err
+				}
+				layers = append(layers, imagegen.LayerInfo{
+					Digest:    qbiasLayer.Digest,
+					Size:      qbiasLayer.Size,
+					MediaType: qbiasLayer.MediaType,
+					Name:      name + "_qbias", // Add _qbias suffix
+				})
+			}
+
+			return layers, nil
+		}
+
+		// Non-quantized path: just create a single layer
+		layer, err := server.NewLayer(r, server.MediaTypeImageTensor)
+		if err != nil {
+			return nil, err
+		}
+
+		return []imagegen.LayerInfo{
+			{
+				Digest:    layer.Digest,
+				Size:      layer.Size,
+				MediaType: layer.MediaType,
+				Name:      name,
+			},
+		}, nil
+	}
+
+	// Create manifest writer callback
+	writeManifest := func(modelName string, config imagegen.LayerInfo, layers []imagegen.LayerInfo) error {
+		name := model.ParseName(modelName)
+		if !name.IsValid() {
+			return fmt.Errorf("invalid model name: %s", modelName)
+		}
+
+		// Create a proper config blob with version requirement
+		configData := model.ConfigV2{
+			ModelFormat:  "safetensors",
+			Capabilities: []string{"image"},
+			Requires:     MinOllamaVersion,
+		}
+		configJSON, err := json.Marshal(configData)
+		if err != nil {
+			return fmt.Errorf("failed to marshal config: %w", err)
+		}
+
+		// Create config layer blob
+		configLayer, err := server.NewLayer(bytes.NewReader(configJSON), "application/vnd.docker.container.image.v1+json")
+		if err != nil {
+			return fmt.Errorf("failed to create config layer: %w", err)
+		}
+
+		// Convert LayerInfo to server.Layer (include the original model_index.json in layers)
+		serverLayers := make([]server.Layer, len(layers))
+		for i, l := range layers {
+			serverLayers[i] = server.Layer{
+				MediaType: l.MediaType,
+				Digest:    l.Digest,
+				Size:      l.Size,
+				Name:      l.Name,
+			}
+		}
+
+		return server.WriteManifest(name, configLayer, serverLayers)
+	}
+
+	// Progress callback
+	progressFn := func(msg string) {
+		spinner.Stop()
+		status = msg
+		spinner = progress.NewSpinner(status)
+		p.Add("imagegen", spinner)
+	}
+
+	err := imagegen.CreateModel(modelName, modelDir, quantize, createLayer, createTensorLayer, writeManifest, progressFn)
+	spinner.Stop()
+	if err != nil {
+		return err
+	}
+
+	fmt.Printf("Created image generation model '%s'\n", modelName)
+	return nil
+}

x/imagegen/client/quantize.go

@@ -11,11 +11,10 @@ import (
 	"github.com/ollama/ollama/x/imagegen/mlx"
 )
 
-// quantizeTensor loads a tensor from safetensors format, quantizes it,
+// quantizeTensor loads a tensor from safetensors format, quantizes it to affine int8,
 // and returns safetensors data for the quantized weights, scales, and biases.
-// Supported quantization types: "fp8" (affine 8-bit)
 // Uses MLX's native SaveSafetensors to ensure correct dtype handling (especially uint32 for quantized weights).
-func quantizeTensor(r io.Reader, name, dtype string, shape []int32, quantize string) (qweightData, scalesData, qbiasData []byte, qweightShape, scalesShape, qbiasShape []int32, err error) {
+func quantizeTensor(r io.Reader, name, dtype string, shape []int32) (qweightData, scalesData, qbiasData []byte, qweightShape, scalesShape, qbiasShape []int32, err error) {
 	tmpDir := ensureTempDir()
 
 	// Read safetensors data to a temp file (LoadSafetensorsNative needs a path)
@@ -51,18 +50,9 @@ func quantizeTensor(r io.Reader, name, dtype string, shape []int32, quantize str
 		mlx.Eval(arr)
 	}
 
-	// Quantize based on quantization type
-	var qweight, scales, qbiases *mlx.Array
-	switch quantize {
-	case "fp4":
-		// affine mode: group_size=32, bits=4
-		qweight, scales, qbiases = mlx.Quantize(arr, 32, 4, "affine")
-	case "fp8":
-		// affine mode: group_size=32, bits=8
-		qweight, scales, qbiases = mlx.Quantize(arr, 32, 8, "affine")
-	default:
-		return nil, nil, nil, nil, nil, nil, fmt.Errorf("unsupported quantization type: %s", quantize)
-	}
+	// Quantize with affine mode: group_size=32, bits=8
+	// Note: mxfp8 mode doesn't have matmul kernels in MLX, affine mode does
+	qweight, scales, qbiases := mlx.Quantize(arr, 32, 8, "affine")
 
 	// Eval and make contiguous for data access
 	qweight = mlx.Contiguous(qweight)

x/imagegen/client/quantize_stub.go

@@ -8,7 +8,7 @@ import (
 )
 
 // quantizeTensor is not available without MLX
-func quantizeTensor(r io.Reader, name, dtype string, shape []int32, quantize string) (qweightData, scalesData, qbiasData []byte, qweightShape, scalesShape, qbiasShape []int32, err error) {
+func quantizeTensor(r io.Reader, name, dtype string, shape []int32) (qweightData, scalesData, qbiasData []byte, qweightShape, scalesShape, qbiasShape []int32, err error) {
 	return nil, nil, nil, nil, nil, nil, fmt.Errorf("quantization requires MLX support (build with mlx tag)")
 }
 

x/imagegen/cmd/engine/generate.go

@@ -65,12 +65,12 @@ func (s *utf8Streamer) Flush() string {
 	return result
 }
 
+func init() {
+	generationStream = mlx.NewStream()
+}
+
 // withStream runs fn with the generation stream as default
 func withStream(fn func()) {
-	// Lazy initialization of generationStream
-	if generationStream == nil {
-		generationStream = mlx.NewStream()
-	}
 	orig := mlx.GetDefaultStream()
 	mlx.SetDefaultStream(generationStream)
 	fn()

x/imagegen/cmd/engine/main.go

@@ -7,17 +7,12 @@ import (
 	"encoding/json"
 	"flag"
 	"fmt"
-	"image"
-	_ "image/jpeg"
-	_ "image/png"
 	"log"
 	"os"
 	"path/filepath"
 	"runtime/pprof"
 
-	"github.com/ollama/ollama/x/imagegen"
 	"github.com/ollama/ollama/x/imagegen/mlx"
-	"github.com/ollama/ollama/x/imagegen/models/flux2"
 	"github.com/ollama/ollama/x/imagegen/models/gemma3"
 	"github.com/ollama/ollama/x/imagegen/models/gpt_oss"
 	"github.com/ollama/ollama/x/imagegen/models/llama"
@@ -51,9 +46,9 @@ func main() {
 	imagePath := flag.String("image", "", "Image path for multimodal models")
 
 	// Image generation params
-	width := flag.Int("width", 0, "Image width (0 = auto from input or 1024)")
-	height := flag.Int("height", 0, "Image height (0 = auto from input or 1024)")
-	steps := flag.Int("steps", 0, "Denoising steps (0 = model default)")
+	width := flag.Int("width", 1024, "Image width")
+	height := flag.Int("height", 1024, "Image height")
+	steps := flag.Int("steps", 9, "Denoising steps")
 	seed := flag.Int64("seed", 42, "Random seed")
 	out := flag.String("output", "output.png", "Output path")
 
@@ -66,7 +61,6 @@ func main() {
 
 	// Legacy mode flags
 	zimageFlag := flag.Bool("zimage", false, "Z-Image generation")
-	flux2Flag := flag.Bool("flux2", false, "FLUX.2 Klein generation")
 	qwenImage := flag.Bool("qwen-image", false, "Qwen-Image text-to-image generation")
 	qwenImageEdit := flag.Bool("qwen-image-edit", false, "Qwen-Image-Edit image editing")
 	var inputImages stringSlice
@@ -84,11 +78,6 @@ func main() {
 		return
 	}
 
-	// Check if MLX initialized successfully
-	if !mlx.IsMLXAvailable() {
-		log.Fatalf("MLX initialization failed: %v", mlx.GetMLXInitError())
-	}
-
 	// CPU profiling
 	if *cpuProfile != "" {
 		f, err := os.Create(*cpuProfile)
@@ -128,44 +117,6 @@ func main() {
 		if err == nil {
 			err = saveImageArray(img, *out)
 		}
-	case *flux2Flag:
-		m := &flux2.Model{}
-		if loadErr := m.Load(*modelPath); loadErr != nil {
-			log.Fatal(loadErr)
-		}
-		// Load input images with EXIF orientation correction
-		var loadedImages []image.Image
-		for _, path := range inputImages {
-			img, loadErr := loadImageWithEXIF(path)
-			if loadErr != nil {
-				log.Fatalf("Failed to load image %s: %v", path, loadErr)
-			}
-			loadedImages = append(loadedImages, img)
-		}
-		// When input images provided and user didn't override dimensions, use 0 to match input
-		fluxWidth := int32(*width)
-		fluxHeight := int32(*height)
-		if len(loadedImages) > 0 && *width == 0 && *height == 0 {
-			// Both unset, will auto-detect from input
-		} else if len(loadedImages) > 0 && *width == 0 {
-			fluxWidth = 0 // Compute from height + aspect ratio
-		} else if len(loadedImages) > 0 && *height == 0 {
-			fluxHeight = 0 // Compute from width + aspect ratio
-		}
-		var img *mlx.Array
-		img, err = m.GenerateFromConfig(context.Background(), &flux2.GenerateConfig{
-			Prompt:        *prompt,
-			Width:         fluxWidth,
-			Height:        fluxHeight,
-			Steps:         *steps,
-			GuidanceScale: float32(*cfgScale),
-			Seed:          *seed,
-			CapturePath:   *gpuCapture,
-			InputImages:   loadedImages,
-		})
-		if err == nil {
-			err = saveImageArray(img, *out)
-		}
 	case *qwenImage:
 		m, loadErr := qwen_image.LoadPersistent(*modelPath)
 		if loadErr != nil {
@@ -320,8 +271,6 @@ func detectModelKind(modelPath string) (string, error) {
 			switch index.ClassName {
 			case "FluxPipeline", "ZImagePipeline":
 				return "zimage", nil
-			case "Flux2KleinPipeline":
-				return "flux2", nil
 			}
 		}
 		return "zimage", nil
@@ -342,12 +291,3 @@ func detectModelKind(modelPath string) (string, error) {
 
 	return cfg.ModelType, nil
 }
-
-// loadImageWithEXIF loads an image from a file path with EXIF orientation correction.
-func loadImageWithEXIF(path string) (image.Image, error) {
-	data, err := os.ReadFile(path)
-	if err != nil {
-		return nil, fmt.Errorf("read file: %w", err)
-	}
-	return imagegen.DecodeImage(data)
-}

x/imagegen/create.go

@@ -1,4 +1,4 @@
-package create
+package imagegen
 
 import (
 	"bytes"
@@ -12,24 +12,40 @@ import (
 	"github.com/ollama/ollama/x/imagegen/safetensors"
 )
 
-// CreateImageGenModel imports an image generation model from a directory.
-// Stores each tensor as a separate blob for fine-grained deduplication.
-// If quantize is specified, linear weights in transformer/text_encoder are quantized.
-// Supported quantization types: fp8 (or empty for no quantization).
-// Layer creation and manifest writing are done via callbacks to avoid import cycles.
-func CreateImageGenModel(modelName, modelDir, quantize string, createLayer LayerCreator, createTensorLayer QuantizingTensorLayerCreator, writeManifest ManifestWriter, fn func(status string)) error {
-	// Validate quantization type
-	switch quantize {
-	case "", "fp4", "fp8":
-		// valid
-	default:
-		return fmt.Errorf("unsupported quantization type %q: supported types are fp4, fp8", quantize)
-	}
+// IsTensorModelDir checks if the directory contains a tensor model
+// by looking for model_index.json, which is the standard diffusers pipeline config.
+func IsTensorModelDir(dir string) bool {
+	_, err := os.Stat(filepath.Join(dir, "model_index.json"))
+	return err == nil
+}
 
+// LayerInfo holds metadata for a created layer.
+type LayerInfo struct {
+	Digest    string
+	Size      int64
+	MediaType string
+	Name      string // Path-style name: "component/tensor" or "path/to/config.json"
+}
+
+// LayerCreator is called to create a blob layer.
+// name is the path-style name (e.g., "tokenizer/tokenizer.json")
+type LayerCreator func(r io.Reader, mediaType, name string) (LayerInfo, error)
+
+// TensorLayerCreator creates a tensor blob layer with metadata.
+// name is the path-style name including component (e.g., "text_encoder/model.embed_tokens.weight")
+type TensorLayerCreator func(r io.Reader, name, dtype string, shape []int32) (LayerInfo, error)
+
+// ManifestWriter writes the manifest file.
+type ManifestWriter func(modelName string, config LayerInfo, layers []LayerInfo) error
+
+// CreateModel imports an image generation model from a directory.
+// Stores each tensor as a separate blob for fine-grained deduplication.
+// If quantize is "fp8", linear weights in transformer/text_encoder are quantized to mxfp8 format.
+// Layer creation and manifest writing are done via callbacks to avoid import cycles.
+func CreateModel(modelName, modelDir, quantize string, createLayer LayerCreator, createTensorLayer QuantizingTensorLayerCreator, writeManifest ManifestWriter, fn func(status string)) error {
 	var layers []LayerInfo
 	var configLayer LayerInfo
 	var totalParams int64 // Count parameters from original tensor shapes
-	var torchDtype string // Read from component config for quantization display
 
 	// Components to process - extract individual tensors from each
 	components := []string{"text_encoder", "transformer", "vae"}
@@ -61,8 +77,8 @@ func CreateImageGenModel(modelName, modelDir, quantize string, createLayer Layer
 
 			tensorNames := extractor.ListTensors()
 			quantizeMsg := ""
-			if quantize != "" && component != "vae" {
-				quantizeMsg = ", quantizing to " + quantize
+			if quantize == "fp8" && component != "vae" {
+				quantizeMsg = ", quantizing to fp8"
 			}
 			fn(fmt.Sprintf("importing %s/%s (%d tensors%s)", component, entry.Name(), len(tensorNames), quantizeMsg))
 
@@ -87,14 +103,11 @@ func CreateImageGenModel(modelName, modelDir, quantize string, createLayer Layer
 				// Use path-style name: "component/tensor_name"
 				fullName := component + "/" + tensorName
 
-				// Determine quantization type for this tensor (empty string if not quantizing)
-				quantizeType := ""
-				if quantize != "" && ShouldQuantize(tensorName, component) && canQuantizeShape(td.Shape) {
-					quantizeType = quantize
-				}
+				// Determine if this tensor should be quantized
+				doQuantize := quantize == "fp8" && ShouldQuantize(tensorName, component)
 
 				// createTensorLayer returns multiple layers if quantizing (weight + scales)
-				newLayers, err := createTensorLayer(td.SafetensorsReader(), fullName, td.Dtype, td.Shape, quantizeType)
+				newLayers, err := createTensorLayer(td.SafetensorsReader(), fullName, td.Dtype, td.Shape, doQuantize)
 				if err != nil {
 					extractor.Close()
 					return fmt.Errorf("failed to create layer for %s: %w", fullName, err)
@@ -106,19 +119,6 @@ func CreateImageGenModel(modelName, modelDir, quantize string, createLayer Layer
 		}
 	}
 
-	// Read torch_dtype from text_encoder config for quantization display
-	if torchDtype == "" {
-		textEncoderConfig := filepath.Join(modelDir, "text_encoder/config.json")
-		if data, err := os.ReadFile(textEncoderConfig); err == nil {
-			var cfg struct {
-				TorchDtype string `json:"torch_dtype"`
-			}
-			if json.Unmarshal(data, &cfg) == nil && cfg.TorchDtype != "" {
-				torchDtype = cfg.TorchDtype
-			}
-		}
-	}
-
 	// Import config files
 	configFiles := []string{
 		"model_index.json",
@@ -164,11 +164,11 @@ func CreateImageGenModel(modelName, modelDir, quantize string, createLayer Layer
 			// Add parameter count (counted from tensor shapes during import)
 			cfg["parameter_count"] = totalParams
 
-			// Add quantization info - use quantize type if set, otherwise torch_dtype
-			if quantize != "" {
-				cfg["quantization"] = strings.ToUpper(quantize)
+			// Add quantization info
+			if quantize == "fp8" {
+				cfg["quantization"] = "FP8"
 			} else {
-				cfg["quantization"] = torchDtype
+				cfg["quantization"] = "BF16"
 			}
 
 			data, err = json.MarshalIndent(cfg, "", "    ")
@@ -211,12 +211,3 @@ func CreateImageGenModel(modelName, modelDir, quantize string, createLayer Layer
 	fn(fmt.Sprintf("successfully imported %s with %d layers", modelName, len(layers)))
 	return nil
 }
-
-// canQuantizeShape returns true if a tensor shape is compatible with MLX quantization.
-// MLX requires the last dimension to be divisible by the group size (32).
-func canQuantizeShape(shape []int32) bool {
-	if len(shape) < 2 {
-		return false
-	}
-	return shape[len(shape)-1]%32 == 0
-}

x/imagegen/image.go

@@ -7,7 +7,6 @@ import (
 	"encoding/base64"
 	"fmt"
 	"image"
-	_ "image/jpeg"
 	"image/png"
 	"os"
 	"path/filepath"
@@ -109,160 +108,3 @@ func clampF(v, min, max float32) float32 {
 	}
 	return v
 }
-
-// DecodeImage decodes image bytes with EXIF orientation applied.
-func DecodeImage(data []byte) (image.Image, error) {
-	orientation := readJPEGOrientation(data)
-
-	img, _, err := image.Decode(bytes.NewReader(data))
-	if err != nil {
-		return nil, err
-	}
-
-	return applyOrientation(img, orientation), nil
-}
-
-// readJPEGOrientation extracts EXIF orientation from JPEG bytes.
-// Returns 1 (normal) for non-JPEG or if orientation not found.
-func readJPEGOrientation(data []byte) int {
-	if len(data) < 2 || data[0] != 0xFF || data[1] != 0xD8 {
-		return 1 // Not JPEG
-	}
-
-	r := bytes.NewReader(data[2:])
-	for {
-		var marker [2]byte
-		if _, err := r.Read(marker[:]); err != nil || marker[0] != 0xFF {
-			return 1
-		}
-
-		if marker[1] == 0xE1 { // APP1 (EXIF)
-			var lenBytes [2]byte
-			if _, err := r.Read(lenBytes[:]); err != nil {
-				return 1
-			}
-			segLen := int(uint16(lenBytes[0])<<8|uint16(lenBytes[1])) - 2
-			if segLen < 14 {
-				r.Seek(int64(segLen), 1)
-				continue
-			}
-			seg := make([]byte, segLen)
-			if _, err := r.Read(seg); err != nil {
-				return 1
-			}
-			if string(seg[:4]) == "Exif" && seg[4] == 0 && seg[5] == 0 {
-				return parseTIFFOrientation(seg[6:])
-			}
-			continue
-		}
-
-		if marker[1] == 0xD9 || marker[1] == 0xDA {
-			return 1 // EOI or SOS
-		}
-		if marker[1] >= 0xD0 && marker[1] <= 0xD7 {
-			continue // RST markers
-		}
-
-		var lenBytes [2]byte
-		if _, err := r.Read(lenBytes[:]); err != nil {
-			return 1
-		}
-		segLen := int(uint16(lenBytes[0])<<8|uint16(lenBytes[1])) - 2
-		if segLen > 0 {
-			r.Seek(int64(segLen), 1)
-		}
-	}
-}
-
-func parseTIFFOrientation(tiff []byte) int {
-	if len(tiff) < 8 {
-		return 1
-	}
-
-	var big bool
-	switch string(tiff[:2]) {
-	case "MM":
-		big = true
-	case "II":
-		big = false
-	default:
-		return 1
-	}
-
-	u16 := func(b []byte) uint16 {
-		if big {
-			return uint16(b[0])<<8 | uint16(b[1])
-		}
-		return uint16(b[1])<<8 | uint16(b[0])
-	}
-	u32 := func(b []byte) uint32 {
-		if big {
-			return uint32(b[0])<<24 | uint32(b[1])<<16 | uint32(b[2])<<8 | uint32(b[3])
-		}
-		return uint32(b[3])<<24 | uint32(b[2])<<16 | uint32(b[1])<<8 | uint32(b[0])
-	}
-
-	if u16(tiff[2:4]) != 42 {
-		return 1
-	}
-
-	ifdOffset := u32(tiff[4:8])
-	if int(ifdOffset)+2 > len(tiff) {
-		return 1
-	}
-
-	numEntries := u16(tiff[ifdOffset : ifdOffset+2])
-	for i := range int(numEntries) {
-		offset := ifdOffset + 2 + uint32(i)*12
-		if int(offset)+12 > len(tiff) {
-			break
-		}
-		if u16(tiff[offset:offset+2]) == 0x0112 { // Orientation tag
-			o := int(u16(tiff[offset+8 : offset+10]))
-			if o >= 1 && o <= 8 {
-				return o
-			}
-			return 1
-		}
-	}
-	return 1
-}
-
-func applyOrientation(img image.Image, orientation int) image.Image {
-	if orientation <= 1 || orientation > 8 {
-		return img
-	}
-
-	bounds := img.Bounds()
-	w, h := bounds.Dx(), bounds.Dy()
-
-	outW, outH := w, h
-	if orientation >= 5 {
-		outW, outH = h, w
-	}
-
-	out := image.NewRGBA(image.Rect(0, 0, outW, outH))
-	for y := range h {
-		for x := range w {
-			var dx, dy int
-			switch orientation {
-			case 2:
-				dx, dy = w-1-x, y
-			case 3:
-				dx, dy = w-1-x, h-1-y
-			case 4:
-				dx, dy = x, h-1-y
-			case 5:
-				dx, dy = y, x
-			case 6:
-				dx, dy = h-1-y, x
-			case 7:
-				dx, dy = h-1-y, w-1-x
-			case 8:
-				dx, dy = y, w-1-x
-			}
-			out.Set(dx, dy, img.At(x+bounds.Min.X, y+bounds.Min.Y))
-		}
-	}
-	return out
-}

x/imagegen/manifest.go

@@ -175,63 +175,3 @@ func (m *ModelManifest) HasTensorLayers() bool {
 	}
 	return false
 }
-
-// ModelInfo contains metadata about an image generation model.
-type ModelInfo struct {
-	Architecture   string
-	ParameterCount int64
-	Quantization   string
-}
-
-// GetModelInfo returns metadata about an image generation model.
-func GetModelInfo(modelName string) (*ModelInfo, error) {
-	manifest, err := LoadManifest(modelName)
-	if err != nil {
-		return nil, fmt.Errorf("failed to load manifest: %w", err)
-	}
-
-	info := &ModelInfo{}
-
-	// Read model_index.json for architecture, parameter count, and quantization
-	if data, err := manifest.ReadConfig("model_index.json"); err == nil {
-		var index struct {
-			Architecture   string `json:"architecture"`
-			ParameterCount int64  `json:"parameter_count"`
-			Quantization   string `json:"quantization"`
-		}
-		if json.Unmarshal(data, &index) == nil {
-			info.Architecture = index.Architecture
-			info.ParameterCount = index.ParameterCount
-			info.Quantization = index.Quantization
-		}
-	}
-
-	// Fallback: detect quantization from tensor names if not in config
-	if info.Quantization == "" {
-		for _, layer := range manifest.Manifest.Layers {
-			if strings.HasSuffix(layer.Name, ".weight_scale") {
-				info.Quantization = "FP8"
-				break
-			}
-		}
-		if info.Quantization == "" {
-			info.Quantization = "BF16"
-		}
-	}
-
-	// Fallback: estimate parameter count if not in config
-	if info.ParameterCount == 0 {
-		var totalSize int64
-		for _, layer := range manifest.Manifest.Layers {
-			if layer.MediaType == "application/vnd.ollama.image.tensor" {
-				if !strings.HasSuffix(layer.Name, "_scale") && !strings.HasSuffix(layer.Name, "_qbias") {
-					totalSize += layer.Size
-				}
-			}
-		}
-		// Assume BF16 (2 bytes/param) as rough estimate
-		info.ParameterCount = totalSize / 2
-	}
-
-	return info, nil
-}

x/imagegen/memory.go

@@ -24,8 +24,9 @@ var SupportedBackends = []string{"metal", "cuda", "cpu"}
 
 // modelVRAMEstimates maps pipeline class names to their estimated VRAM requirements.
 var modelVRAMEstimates = map[string]uint64{
-	"ZImagePipeline": 21 * GB, // ~21GB for Z-Image (text encoder + transformer + VAE)
-	"FluxPipeline":   20 * GB, // ~20GB for Flux
+	"ZImagePipeline":    21 * GB, // ~21GB for Z-Image (text encoder + transformer + VAE)
+	"FluxPipeline":      21 * GB, // ~21GB for Flux (same architecture)
+	"QwenImagePipeline": 80 * GB, // TODO: verify actual requirements, using conservative estimate for now
 }
 
 // CheckPlatformSupport validates that image generation is supported on the current platform.
@@ -71,38 +72,31 @@ func ResolveModelName(modelName string) string {
 // EstimateVRAM returns the estimated VRAM needed for an image generation model.
 // Returns a conservative default of 21GB if the model type cannot be determined.
 func EstimateVRAM(modelName string) uint64 {
-	className := DetectModelType(modelName)
-	if estimate, ok := modelVRAMEstimates[className]; ok {
+	manifest, err := LoadManifest(modelName)
+	if err != nil {
+		return 21 * GB
+	}
+
+	data, err := manifest.ReadConfig("model_index.json")
+	if err != nil {
+		return 21 * GB
+	}
+
+	// Parse just the class name
+	var index struct {
+		ClassName string `json:"_class_name"`
+	}
+	if err := json.Unmarshal(data, &index); err != nil {
+		return 21 * GB
+	}
+
+	if estimate, ok := modelVRAMEstimates[index.ClassName]; ok {
 		return estimate
 	}
 	return 21 * GB
 }
 
-// DetectModelType reads model_index.json and returns the model type.
-// Checks both "architecture" (Ollama format) and "_class_name" (diffusers format).
-// Returns empty string if detection fails.
-func DetectModelType(modelName string) string {
-	manifest, err := LoadManifest(modelName)
-	if err != nil {
-		return ""
-	}
-
-	data, err := manifest.ReadConfig("model_index.json")
-	if err != nil {
-		return ""
-	}
-
-	var index struct {
-		Architecture string `json:"architecture"`
-		ClassName    string `json:"_class_name"`
-	}
-	if err := json.Unmarshal(data, &index); err != nil {
-		return ""
-	}
-
-	// Prefer architecture (Ollama format), fall back to _class_name (diffusers)
-	if index.Architecture != "" {
-		return index.Architecture
-	}
-	return index.ClassName
+// HasTensorLayers checks if the given model has tensor layers.
+func HasTensorLayers(modelName string) bool {
+	return ResolveModelName(modelName) != ""
 }

x/imagegen/memory_test.go

@@ -72,8 +72,9 @@ func TestCheckMemoryRequirements(t *testing.T) {
 func TestModelVRAMEstimates(t *testing.T) {
 	// Verify the VRAM estimates map has expected entries
 	expected := map[string]uint64{
-		"ZImagePipeline": 21 * GB,
-		"FluxPipeline":   20 * GB,
+		"ZImagePipeline":    21 * GB,
+		"FluxPipeline":      21 * GB,
+		"QwenImagePipeline": 80 * GB,
 	}
 
 	for name, expectedVRAM := range expected {
@@ -93,6 +94,13 @@ func TestEstimateVRAMDefault(t *testing.T) {
 	}
 }
 
+func TestHasTensorLayers(t *testing.T) {
+	// Non-existent model should return false
+	if HasTensorLayers("nonexistent-model") {
+		t.Error("HasTensorLayers() should return false for non-existent model")
+	}
+}
+
 func TestResolveModelName(t *testing.T) {
 	// Non-existent model should return empty string
 	result := ResolveModelName("nonexistent-model")

x/imagegen/mlx/compile.go

@@ -3,7 +3,7 @@
 package mlx
 
 /*
-#include "mlx.h"
+#include "mlx/c/mlx.h"
 #include <stdlib.h>
 
 // Forward declaration for Go callback

x/imagegen/mlx/doc.go

@@ -1,6 +0,0 @@
-//go:build mlx
-
-// Package mlx provides Go bindings for the MLX-C library with dynamic loading support.
-//
-//go:generate go run generate_wrappers.go ../../../build/_deps/mlx-c-src/mlx/c mlx.h mlx.c
-package mlx

x/imagegen/mlx/generate_wrappers.go

@@ -1,439 +0,0 @@
-//go:build ignore
-
-// This tool generates MLX-C dynamic loading wrappers.
-// Usage: go run generate_wrappers.go <mlx-c-include-dir> <output-header> [output-impl]
-package main
-
-import (
-	"bytes"
-	"flag"
-	"fmt"
-	"io/fs"
-	"os"
-	"path/filepath"
-	"regexp"
-	"strings"
-)
-
-type Function struct {
-	Name          string
-	ReturnType    string
-	Params        string
-	ParamNames    []string
-	NeedsARM64Guard bool
-}
-
-func findHeaders(directory string) ([]string, error) {
-	var headers []string
-	err := filepath.WalkDir(directory, func(path string, d fs.DirEntry, err error) error {
-		if err != nil {
-			return err
-		}
-		if !d.IsDir() && strings.HasSuffix(path, ".h") {
-			headers = append(headers, path)
-		}
-		return nil
-	})
-	return headers, err
-}
-
-func cleanContent(content string) string {
-	// Remove single-line comments
-	re := regexp.MustCompile(`//.*?\n`)
-	content = re.ReplaceAllString(content, "\n")
-
-	// Remove multi-line comments
-	re = regexp.MustCompile(`/\*.*?\*/`)
-	content = re.ReplaceAllString(content, "")
-
-	// Remove preprocessor directives (lines starting with #) - use multiline mode
-	re = regexp.MustCompile(`(?m)^\s*#.*?$`)
-	content = re.ReplaceAllString(content, "")
-
-	// Remove extern "C" { and } blocks more conservatively
-	// Only remove the extern "C" { line, not the content inside
-	re = regexp.MustCompile(`extern\s+"C"\s*\{\s*?\n`)
-	content = re.ReplaceAllString(content, "\n")
-	// Remove standalone closing braces that are not part of function declarations
-	re = regexp.MustCompile(`\n\s*\}\s*\n`)
-	content = re.ReplaceAllString(content, "\n")
-
-	// Collapse whitespace and newlines
-	re = regexp.MustCompile(`\s+`)
-	content = re.ReplaceAllString(content, " ")
-
-	return content
-}
-
-func extractParamNames(params string) []string {
-	if params == "" || strings.TrimSpace(params) == "void" {
-		return []string{}
-	}
-
-	var names []string
-
-	// Split by comma, but respect parentheses (for function pointers)
-	parts := splitParams(params)
-
-	// Remove array brackets
-	arrayBrackets := regexp.MustCompile(`\[.*?\]`)
-
-	// Function pointer pattern
-	funcPtrPattern := regexp.MustCompile(`\(\s*\*\s*(\w+)\s*\)`)
-
-	// Type keywords to skip
-	typeKeywords := map[string]bool{
-		"const":     true,
-		"struct":    true,
-		"unsigned":  true,
-		"signed":    true,
-		"long":      true,
-		"short":     true,
-		"int":       true,
-		"char":      true,
-		"float":     true,
-		"double":    true,
-		"void":      true,
-		"size_t":    true,
-		"uint8_t":   true,
-		"uint16_t":  true,
-		"uint32_t":  true,
-		"uint64_t":  true,
-		"int8_t":    true,
-		"int16_t":   true,
-		"int32_t":   true,
-		"int64_t":   true,
-		"intptr_t":  true,
-		"uintptr_t": true,
-	}
-
-	for _, part := range parts {
-		if part == "" {
-			continue
-		}
-
-		// Remove array brackets
-		part = arrayBrackets.ReplaceAllString(part, "")
-
-		// For function pointers like "void (*callback)(int)"
-		if matches := funcPtrPattern.FindStringSubmatch(part); len(matches) > 1 {
-			names = append(names, matches[1])
-			continue
-		}
-
-		// Regular parameter: last identifier
-		tokens := regexp.MustCompile(`\w+`).FindAllString(part, -1)
-		if len(tokens) > 0 {
-			// The last token is usually the parameter name
-			// Skip type keywords
-			for i := len(tokens) - 1; i >= 0; i-- {
-				if !typeKeywords[tokens[i]] {
-					names = append(names, tokens[i])
-					break
-				}
-			}
-		}
-	}
-
-	return names
-}
-
-func splitParams(params string) []string {
-	var parts []string
-	var current bytes.Buffer
-	depth := 0
-
-	for _, char := range params + "," {
-		switch char {
-		case '(':
-			depth++
-			current.WriteRune(char)
-		case ')':
-			depth--
-			current.WriteRune(char)
-		case ',':
-			if depth == 0 {
-				parts = append(parts, strings.TrimSpace(current.String()))
-				current.Reset()
-			} else {
-				current.WriteRune(char)
-			}
-		default:
-			current.WriteRune(char)
-		}
-	}
-
-	return parts
-}
-
-func parseFunctions(content string) []Function {
-	var functions []Function
-
-	// Match function declarations: return_type function_name(params);
-	// Matches both mlx_* and _mlx_* functions
-	pattern := regexp.MustCompile(`\b((?:const\s+)?(?:struct\s+)?[\w\s]+?[\*\s]*)\s+(_?mlx_\w+)\s*\(([^)]*(?:\([^)]*\)[^)]*)*)\)\s*;`)
-
-	matches := pattern.FindAllStringSubmatch(content, -1)
-	for _, match := range matches {
-		returnType := strings.TrimSpace(match[1])
-		funcName := strings.TrimSpace(match[2])
-		params := strings.TrimSpace(match[3])
-
-		// Skip if this looks like a variable declaration
-		if params == "" || strings.Contains(params, "{") {
-			continue
-		}
-
-		// Clean up return type
-		returnType = strings.Join(strings.Fields(returnType), " ")
-
-		// Extract parameter names
-		paramNames := extractParamNames(params)
-
-		// Check if ARM64 guard is needed
-		needsGuard := needsARM64Guard(funcName, returnType, params)
-
-		functions = append(functions, Function{
-			Name:           funcName,
-			ReturnType:     returnType,
-			Params:         params,
-			ParamNames:     paramNames,
-			NeedsARM64Guard: needsGuard,
-		})
-	}
-
-	return functions
-}
-
-func needsARM64Guard(name, retType, params string) bool {
-	return strings.Contains(name, "float16") ||
-		strings.Contains(name, "bfloat16") ||
-		strings.Contains(retType, "float16_t") ||
-		strings.Contains(retType, "bfloat16_t") ||
-		strings.Contains(params, "float16_t") ||
-		strings.Contains(params, "bfloat16_t")
-}
-
-func generateWrapperFiles(functions []Function, headerPath, implPath string) error {
-	// Generate header file
-	var headerBuf bytes.Buffer
-
-	headerBuf.WriteString("// AUTO-GENERATED by generate_wrappers.go - DO NOT EDIT\n")
-	headerBuf.WriteString("// This file provides wrapper declarations for MLX-C functions that use dlopen/dlsym\n")
-	headerBuf.WriteString("//\n")
-	headerBuf.WriteString("// Strategy: Include MLX-C headers for type definitions, then provide wrapper\n")
-	headerBuf.WriteString("// functions that shadow the originals, allowing Go code to call them directly (e.g., C.mlx_add).\n")
-	headerBuf.WriteString("// Function pointers are defined in mlx.c (single compilation unit).\n\n")
-	headerBuf.WriteString("#ifndef MLX_WRAPPERS_H\n")
-	headerBuf.WriteString("#define MLX_WRAPPERS_H\n\n")
-
-	headerBuf.WriteString("// Include MLX headers for type definitions and original declarations\n")
-	headerBuf.WriteString("#include \"mlx/c/mlx.h\"\n")
-	headerBuf.WriteString("#include \"mlx_dynamic.h\"\n")
-	headerBuf.WriteString("#include <stdio.h>\n\n")
-
-	// Undef all MLX functions to avoid conflicts
-	headerBuf.WriteString("// Undefine any existing MLX function macros\n")
-	for _, fn := range functions {
-		headerBuf.WriteString(fmt.Sprintf("#undef %s\n", fn.Name))
-	}
-	headerBuf.WriteString("\n")
-
-	// Function pointer extern declarations
-	headerBuf.WriteString("// Function pointer declarations (defined in mlx.c, loaded via dlsym)\n")
-	for _, fn := range functions {
-		if fn.NeedsARM64Guard {
-			headerBuf.WriteString("#if defined(__aarch64__) || defined(_M_ARM64)\n")
-		}
-		headerBuf.WriteString(fmt.Sprintf("extern %s (*%s_ptr)(%s);\n", fn.ReturnType, fn.Name, fn.Params))
-		if fn.NeedsARM64Guard {
-			headerBuf.WriteString("#endif\n")
-		}
-	}
-	headerBuf.WriteString("\n")
-
-	// Initialization function declaration
-	headerBuf.WriteString("// Initialize all function pointers via dlsym (defined in mlx.c)\n")
-	headerBuf.WriteString("int mlx_load_functions(void* handle);\n\n")
-
-	// Wrapper function declarations
-	headerBuf.WriteString("// Wrapper function declarations that call through function pointers\n")
-	headerBuf.WriteString("// Go code calls these directly as C.mlx_* (no #define redirection needed)\n")
-	for _, fn := range functions {
-		if fn.NeedsARM64Guard {
-			headerBuf.WriteString("#if defined(__aarch64__) || defined(_M_ARM64)\n")
-		}
-		headerBuf.WriteString(fmt.Sprintf("%s %s(%s);\n", fn.ReturnType, fn.Name, fn.Params))
-		if fn.NeedsARM64Guard {
-			headerBuf.WriteString("#endif\n")
-		}
-		headerBuf.WriteString("\n")
-	}
-
-	headerBuf.WriteString("#endif // MLX_WRAPPERS_H\n")
-
-	// Write header file
-	if err := os.WriteFile(headerPath, headerBuf.Bytes(), 0644); err != nil {
-		return fmt.Errorf("failed to write header file: %w", err)
-	}
-
-	// Generate implementation file
-	var implBuf bytes.Buffer
-
-	implBuf.WriteString("// AUTO-GENERATED by generate_wrappers.go - DO NOT EDIT\n")
-	implBuf.WriteString("// This file contains the function pointer definitions and initialization\n")
-	implBuf.WriteString("// All function pointers are in a single compilation unit to avoid duplication\n\n")
-
-	implBuf.WriteString("#include \"mlx/c/mlx.h\"\n")
-	implBuf.WriteString("#include \"mlx_dynamic.h\"\n")
-	implBuf.WriteString("#include <stdio.h>\n")
-	implBuf.WriteString("#include <dlfcn.h>\n\n")
-
-	// Function pointer definitions
-	implBuf.WriteString("// Function pointer definitions\n")
-	for _, fn := range functions {
-		if fn.NeedsARM64Guard {
-			implBuf.WriteString("#if defined(__aarch64__) || defined(_M_ARM64)\n")
-		}
-		implBuf.WriteString(fmt.Sprintf("%s (*%s_ptr)(%s) = NULL;\n", fn.ReturnType, fn.Name, fn.Params))
-		if fn.NeedsARM64Guard {
-			implBuf.WriteString("#endif\n")
-		}
-	}
-	implBuf.WriteString("\n")
-
-	// Initialization function
-	implBuf.WriteString("// Initialize all function pointers via dlsym\n")
-	implBuf.WriteString("int mlx_load_functions(void* handle) {\n")
-	implBuf.WriteString("    if (handle == NULL) {\n")
-	implBuf.WriteString("        fprintf(stderr, \"MLX: Invalid library handle\\n\");\n")
-	implBuf.WriteString("        return -1;\n")
-	implBuf.WriteString("    }\n\n")
-
-	for _, fn := range functions {
-		if fn.NeedsARM64Guard {
-			implBuf.WriteString("#if defined(__aarch64__) || defined(_M_ARM64)\n")
-		}
-		implBuf.WriteString(fmt.Sprintf("    %s_ptr = dlsym(handle, \"%s\");\n", fn.Name, fn.Name))
-		implBuf.WriteString(fmt.Sprintf("    if (%s_ptr == NULL) {\n", fn.Name))
-		implBuf.WriteString(fmt.Sprintf("        fprintf(stderr, \"MLX: Failed to load symbol: %s\\n\");\n", fn.Name))
-		implBuf.WriteString("        return -1;\n")
-		implBuf.WriteString("    }\n")
-		if fn.NeedsARM64Guard {
-			implBuf.WriteString("#endif\n")
-		}
-	}
-
-	implBuf.WriteString("    return 0;\n")
-	implBuf.WriteString("}\n\n")
-
-	// Wrapper function implementations
-	implBuf.WriteString("// Wrapper function implementations that call through function pointers\n")
-	for _, fn := range functions {
-		if fn.NeedsARM64Guard {
-			implBuf.WriteString("#if defined(__aarch64__) || defined(_M_ARM64)\n")
-		}
-		implBuf.WriteString(fmt.Sprintf("%s %s(%s) {\n", fn.ReturnType, fn.Name, fn.Params))
-
-		// Call through function pointer
-		if fn.ReturnType != "void" {
-			implBuf.WriteString(fmt.Sprintf("    return %s_ptr(", fn.Name))
-		} else {
-			implBuf.WriteString(fmt.Sprintf("    %s_ptr(", fn.Name))
-		}
-
-		// Pass parameters
-		implBuf.WriteString(strings.Join(fn.ParamNames, ", "))
-		implBuf.WriteString(");\n")
-		implBuf.WriteString("}\n")
-		if fn.NeedsARM64Guard {
-			implBuf.WriteString("#endif\n")
-		}
-		implBuf.WriteString("\n")
-	}
-
-	// Write implementation file
-	if err := os.WriteFile(implPath, implBuf.Bytes(), 0644); err != nil {
-		return fmt.Errorf("failed to write implementation file: %w", err)
-	}
-
-	return nil
-}
-
-func main() {
-	flag.Usage = func() {
-		fmt.Fprintf(flag.CommandLine.Output(), "Usage: go run generate_wrappers.go <mlx-c-include-dir> <output-header> [output-impl]\n")
-		fmt.Fprintf(flag.CommandLine.Output(), "Generate MLX-C dynamic loading wrappers.\n\n")
-		flag.PrintDefaults()
-	}
-	flag.Parse()
-
-	args := flag.Args()
-	if len(args) < 2 {
-		fmt.Fprintf(flag.CommandLine.Output(), "ERROR: Missing required arguments\n\n")
-		flag.Usage()
-		os.Exit(1)
-	}
-
-	headerDir := args[0]
-	outputHeader := args[1]
-	// Default implementation file is same name with .c extension
-	outputImpl := outputHeader
-	if len(args) > 2 {
-		outputImpl = args[2]
-	} else if strings.HasSuffix(outputHeader, ".h") {
-		outputImpl = outputHeader[:len(outputHeader)-2] + ".c"
-	}
-
-	// Check if header directory exists
-	if _, err := os.Stat(headerDir); os.IsNotExist(err) {
-		fmt.Fprintf(os.Stderr, "ERROR: MLX-C headers directory not found at: %s\n\n", headerDir)
-		fmt.Fprintf(os.Stderr, "Please run CMake first to download MLX-C dependencies:\n")
-		fmt.Fprintf(os.Stderr, "  cmake -B build\n\n")
-		fmt.Fprintf(os.Stderr, "The CMake build will download and extract MLX-C headers needed for wrapper generation.\n")
-		os.Exit(1)
-	}
-
-	fmt.Fprintf(os.Stderr, "Parsing MLX-C headers from: %s\n", headerDir)
-
-	// Find all headers
-	headers, err := findHeaders(headerDir)
-	if err != nil {
-		fmt.Fprintf(os.Stderr, "ERROR: Failed to find header files: %v\n", err)
-		os.Exit(1)
-	}
-	fmt.Fprintf(os.Stderr, "Found %d header files\n", len(headers))
-
-	// Parse all headers
-	var allFunctions []Function
-	seen := make(map[string]bool)
-
-	for _, header := range headers {
-		content, err := os.ReadFile(header)
-		if err != nil {
-			fmt.Fprintf(os.Stderr, "Error reading %s: %v\n", header, err)
-			continue
-		}
-
-		cleaned := cleanContent(string(content))
-		functions := parseFunctions(cleaned)
-
-		// Deduplicate
-		for _, fn := range functions {
-			if !seen[fn.Name] {
-				seen[fn.Name] = true
-				allFunctions = append(allFunctions, fn)
-			}
-		}
-	}
-
-	fmt.Fprintf(os.Stderr, "Found %d unique function declarations\n", len(allFunctions))
-
-	// Generate wrapper files
-	if err := generateWrapperFiles(allFunctions, outputHeader, outputImpl); err != nil {
-		fmt.Fprintf(os.Stderr, "ERROR: Failed to generate wrapper files: %v\n", err)
-		os.Exit(1)
-	}
-
-	fmt.Fprintf(os.Stderr, "Generated %s and %s successfully\n", outputHeader, outputImpl)
-}

x/imagegen/mlx/mlx.c

@@ -1,5786 +0,0 @@
-// AUTO-GENERATED by generate_wrappers.go - DO NOT EDIT
-// This file contains the function pointer definitions and initialization
-// All function pointers are in a single compilation unit to avoid duplication
-
-#include "mlx/c/mlx.h"
-#include "mlx_dynamic.h"
-#include <stdio.h>
-#include <dlfcn.h>
-
-// Function pointer definitions
-size_t (*mlx_dtype_size_ptr)(mlx_dtype dtype) = NULL;
-int (*mlx_array_tostring_ptr)(mlx_string* str, const mlx_array arr) = NULL;
-mlx_array (*mlx_array_new_ptr)(void) = NULL;
-int (*mlx_array_free_ptr)(mlx_array arr) = NULL;
-mlx_array (*mlx_array_new_bool_ptr)(bool val) = NULL;
-mlx_array (*mlx_array_new_int_ptr)(int val) = NULL;
-mlx_array (*mlx_array_new_float32_ptr)(float val) = NULL;
-mlx_array (*mlx_array_new_float_ptr)(float val) = NULL;
-mlx_array (*mlx_array_new_float64_ptr)(double val) = NULL;
-mlx_array (*mlx_array_new_double_ptr)(double val) = NULL;
-mlx_array (*mlx_array_new_complex_ptr)(float real_val, float imag_val) = NULL;
-mlx_array (*mlx_array_new_data_ptr)(const void* data, const int* shape, int dim, mlx_dtype dtype) = NULL;
-int (*mlx_array_set_ptr)(mlx_array* arr, const mlx_array src) = NULL;
-int (*mlx_array_set_bool_ptr)(mlx_array* arr, bool val) = NULL;
-int (*mlx_array_set_int_ptr)(mlx_array* arr, int val) = NULL;
-int (*mlx_array_set_float32_ptr)(mlx_array* arr, float val) = NULL;
-int (*mlx_array_set_float_ptr)(mlx_array* arr, float val) = NULL;
-int (*mlx_array_set_float64_ptr)(mlx_array* arr, double val) = NULL;
-int (*mlx_array_set_double_ptr)(mlx_array* arr, double val) = NULL;
-int (*mlx_array_set_complex_ptr)(mlx_array* arr, float real_val, float imag_val) = NULL;
-int (*mlx_array_set_data_ptr)(mlx_array* arr, const void* data, const int* shape, int dim, mlx_dtype dtype) = NULL;
-size_t (*mlx_array_itemsize_ptr)(const mlx_array arr) = NULL;
-size_t (*mlx_array_size_ptr)(const mlx_array arr) = NULL;
-size_t (*mlx_array_nbytes_ptr)(const mlx_array arr) = NULL;
-size_t (*mlx_array_ndim_ptr)(const mlx_array arr) = NULL;
-const int* (*mlx_array_shape_ptr)(const mlx_array arr) = NULL;
-const size_t* (*mlx_array_strides_ptr)(const mlx_array arr) = NULL;
-int (*mlx_array_dim_ptr)(const mlx_array arr, int dim) = NULL;
-mlx_dtype (*mlx_array_dtype_ptr)(const mlx_array arr) = NULL;
-int (*mlx_array_eval_ptr)(mlx_array arr) = NULL;
-int (*mlx_array_item_bool_ptr)(bool* res, const mlx_array arr) = NULL;
-int (*mlx_array_item_uint8_ptr)(uint8_t* res, const mlx_array arr) = NULL;
-int (*mlx_array_item_uint16_ptr)(uint16_t* res, const mlx_array arr) = NULL;
-int (*mlx_array_item_uint32_ptr)(uint32_t* res, const mlx_array arr) = NULL;
-int (*mlx_array_item_uint64_ptr)(uint64_t* res, const mlx_array arr) = NULL;
-int (*mlx_array_item_int8_ptr)(int8_t* res, const mlx_array arr) = NULL;
-int (*mlx_array_item_int16_ptr)(int16_t* res, const mlx_array arr) = NULL;
-int (*mlx_array_item_int32_ptr)(int32_t* res, const mlx_array arr) = NULL;
-int (*mlx_array_item_int64_ptr)(int64_t* res, const mlx_array arr) = NULL;
-int (*mlx_array_item_float32_ptr)(float* res, const mlx_array arr) = NULL;
-int (*mlx_array_item_float64_ptr)(double* res, const mlx_array arr) = NULL;
-int (*mlx_array_item_complex64_ptr)(float _Complex* res, const mlx_array arr) = NULL;
-#if defined(__aarch64__) || defined(_M_ARM64)
-int (*mlx_array_item_float16_ptr)(float16_t* res, const mlx_array arr) = NULL;
-#endif
-#if defined(__aarch64__) || defined(_M_ARM64)
-int (*mlx_array_item_bfloat16_ptr)(bfloat16_t* res, const mlx_array arr) = NULL;
-#endif
-const bool* (*mlx_array_data_bool_ptr)(const mlx_array arr) = NULL;
-const uint8_t* (*mlx_array_data_uint8_ptr)(const mlx_array arr) = NULL;
-const uint16_t* (*mlx_array_data_uint16_ptr)(const mlx_array arr) = NULL;
-const uint32_t* (*mlx_array_data_uint32_ptr)(const mlx_array arr) = NULL;
-const uint64_t* (*mlx_array_data_uint64_ptr)(const mlx_array arr) = NULL;
-const int8_t* (*mlx_array_data_int8_ptr)(const mlx_array arr) = NULL;
-const int16_t* (*mlx_array_data_int16_ptr)(const mlx_array arr) = NULL;
-const int32_t* (*mlx_array_data_int32_ptr)(const mlx_array arr) = NULL;
-const int64_t* (*mlx_array_data_int64_ptr)(const mlx_array arr) = NULL;
-const float* (*mlx_array_data_float32_ptr)(const mlx_array arr) = NULL;
-const double* (*mlx_array_data_float64_ptr)(const mlx_array arr) = NULL;
-const float _Complex* (*mlx_array_data_complex64_ptr)(const mlx_array arr) = NULL;
-#if defined(__aarch64__) || defined(_M_ARM64)
-const float16_t* (*mlx_array_data_float16_ptr)(const mlx_array arr) = NULL;
-#endif
-#if defined(__aarch64__) || defined(_M_ARM64)
-const bfloat16_t* (*mlx_array_data_bfloat16_ptr)(const mlx_array arr) = NULL;
-#endif
-int (*_mlx_array_is_available_ptr)(bool* res, const mlx_array arr) = NULL;
-int (*_mlx_array_wait_ptr)(const mlx_array arr) = NULL;
-int (*_mlx_array_is_contiguous_ptr)(bool* res, const mlx_array arr) = NULL;
-int (*_mlx_array_is_row_contiguous_ptr)(bool* res, const mlx_array arr) = NULL;
-int (*_mlx_array_is_col_contiguous_ptr)(bool* res, const mlx_array arr) = NULL;
-mlx_closure (*mlx_closure_new_ptr)(void) = NULL;
-int (*mlx_closure_free_ptr)(mlx_closure cls) = NULL;
-mlx_closure (*mlx_closure_new_func_ptr)(int (*fun)(mlx_vector_array*, const mlx_vector_array)) = NULL;
-mlx_closure (*mlx_closure_new_func_payload_ptr)(int (*fun)(mlx_vector_array*, const mlx_vector_array, void*), void* payload, void (*dtor)(void*)) = NULL;
-int (*mlx_closure_set_ptr)(mlx_closure* cls, const mlx_closure src) = NULL;
-int (*mlx_closure_apply_ptr)(mlx_vector_array* res, mlx_closure cls, const mlx_vector_array input) = NULL;
-mlx_closure (*mlx_closure_new_unary_ptr)(int (*fun)(mlx_array*, const mlx_array)) = NULL;
-mlx_closure_kwargs (*mlx_closure_kwargs_new_ptr)(void) = NULL;
-int (*mlx_closure_kwargs_free_ptr)(mlx_closure_kwargs cls) = NULL;
-mlx_closure_kwargs (*mlx_closure_kwargs_new_func_ptr)(int (*fun)( mlx_vector_array*, const mlx_vector_array, const mlx_map_string_to_array)) = NULL;
-mlx_closure_kwargs (*mlx_closure_kwargs_new_func_payload_ptr)(int (*fun)( mlx_vector_array*, const mlx_vector_array, const mlx_map_string_to_array, void*), void* payload, void (*dtor)(void*)) = NULL;
-int (*mlx_closure_kwargs_set_ptr)(mlx_closure_kwargs* cls, const mlx_closure_kwargs src) = NULL;
-int (*mlx_closure_kwargs_apply_ptr)(mlx_vector_array* res, mlx_closure_kwargs cls, const mlx_vector_array input_0, const mlx_map_string_to_array input_1) = NULL;
-mlx_closure_value_and_grad (*mlx_closure_value_and_grad_new_ptr)(void) = NULL;
-int (*mlx_closure_value_and_grad_free_ptr)(mlx_closure_value_and_grad cls) = NULL;
-mlx_closure_value_and_grad (*mlx_closure_value_and_grad_new_func_ptr)(int (*fun)(mlx_vector_array*, mlx_vector_array*, const mlx_vector_array)) = NULL;
-mlx_closure_value_and_grad (*mlx_closure_value_and_grad_new_func_payload_ptr)(int (*fun)( mlx_vector_array*, mlx_vector_array*, const mlx_vector_array, void*), void* payload, void (*dtor)(void*)) = NULL;
-int (*mlx_closure_value_and_grad_set_ptr)(mlx_closure_value_and_grad* cls, const mlx_closure_value_and_grad src) = NULL;
-int (*mlx_closure_value_and_grad_apply_ptr)(mlx_vector_array* res_0, mlx_vector_array* res_1, mlx_closure_value_and_grad cls, const mlx_vector_array input) = NULL;
-mlx_closure_custom (*mlx_closure_custom_new_ptr)(void) = NULL;
-int (*mlx_closure_custom_free_ptr)(mlx_closure_custom cls) = NULL;
-mlx_closure_custom (*mlx_closure_custom_new_func_ptr)(int (*fun)( mlx_vector_array*, const mlx_vector_array, const mlx_vector_array, const mlx_vector_array)) = NULL;
-mlx_closure_custom (*mlx_closure_custom_new_func_payload_ptr)(int (*fun)( mlx_vector_array*, const mlx_vector_array, const mlx_vector_array, const mlx_vector_array, void*), void* payload, void (*dtor)(void*)) = NULL;
-int (*mlx_closure_custom_set_ptr)(mlx_closure_custom* cls, const mlx_closure_custom src) = NULL;
-int (*mlx_closure_custom_apply_ptr)(mlx_vector_array* res, mlx_closure_custom cls, const mlx_vector_array input_0, const mlx_vector_array input_1, const mlx_vector_array input_2) = NULL;
-mlx_closure_custom_jvp (*mlx_closure_custom_jvp_new_ptr)(void) = NULL;
-int (*mlx_closure_custom_jvp_free_ptr)(mlx_closure_custom_jvp cls) = NULL;
-mlx_closure_custom_jvp (*mlx_closure_custom_jvp_new_func_ptr)(int (*fun)( mlx_vector_array*, const mlx_vector_array, const mlx_vector_array, const int*, size_t _num)) = NULL;
-mlx_closure_custom_jvp (*mlx_closure_custom_jvp_new_func_payload_ptr)(int (*fun)( mlx_vector_array*, const mlx_vector_array, const mlx_vector_array, const int*, size_t _num, void*), void* payload, void (*dtor)(void*)) = NULL;
-int (*mlx_closure_custom_jvp_set_ptr)(mlx_closure_custom_jvp* cls, const mlx_closure_custom_jvp src) = NULL;
-int (*mlx_closure_custom_jvp_apply_ptr)(mlx_vector_array* res, mlx_closure_custom_jvp cls, const mlx_vector_array input_0, const mlx_vector_array input_1, const int* input_2, size_t input_2_num) = NULL;
-mlx_closure_custom_vmap (*mlx_closure_custom_vmap_new_ptr)(void) = NULL;
-int (*mlx_closure_custom_vmap_free_ptr)(mlx_closure_custom_vmap cls) = NULL;
-mlx_closure_custom_vmap (*mlx_closure_custom_vmap_new_func_ptr)(int (*fun)( mlx_vector_array*, mlx_vector_int*, const mlx_vector_array, const int*, size_t _num)) = NULL;
-mlx_closure_custom_vmap (*mlx_closure_custom_vmap_new_func_payload_ptr)(int (*fun)( mlx_vector_array*, mlx_vector_int*, const mlx_vector_array, const int*, size_t _num, void*), void* payload, void (*dtor)(void*)) = NULL;
-int (*mlx_closure_custom_vmap_set_ptr)(mlx_closure_custom_vmap* cls, const mlx_closure_custom_vmap src) = NULL;
-int (*mlx_closure_custom_vmap_apply_ptr)(mlx_vector_array* res_0, mlx_vector_int* res_1, mlx_closure_custom_vmap cls, const mlx_vector_array input_0, const int* input_1, size_t input_1_num) = NULL;
-int (*mlx_compile_ptr)(mlx_closure* res, const mlx_closure fun, bool shapeless) = NULL;
-int (*mlx_detail_compile_ptr)(mlx_closure* res, const mlx_closure fun, uintptr_t fun_id, bool shapeless, const uint64_t* constants, size_t constants_num) = NULL;
-int (*mlx_detail_compile_clear_cache_ptr)(void) = NULL;
-int (*mlx_detail_compile_erase_ptr)(uintptr_t fun_id) = NULL;
-int (*mlx_disable_compile_ptr)(void) = NULL;
-int (*mlx_enable_compile_ptr)(void) = NULL;
-int (*mlx_set_compile_mode_ptr)(mlx_compile_mode mode) = NULL;
-mlx_device (*mlx_device_new_ptr)(void) = NULL;
-mlx_device (*mlx_device_new_type_ptr)(mlx_device_type type, int index) = NULL;
-int (*mlx_device_free_ptr)(mlx_device dev) = NULL;
-int (*mlx_device_set_ptr)(mlx_device* dev, const mlx_device src) = NULL;
-int (*mlx_device_tostring_ptr)(mlx_string* str, mlx_device dev) = NULL;
-bool (*mlx_device_equal_ptr)(mlx_device lhs, mlx_device rhs) = NULL;
-int (*mlx_device_get_index_ptr)(int* index, mlx_device dev) = NULL;
-int (*mlx_device_get_type_ptr)(mlx_device_type* type, mlx_device dev) = NULL;
-int (*mlx_get_default_device_ptr)(mlx_device* dev) = NULL;
-int (*mlx_set_default_device_ptr)(mlx_device dev) = NULL;
-int (*mlx_distributed_all_gather_ptr)(mlx_array* res, const mlx_array x, const mlx_distributed_group group , const mlx_stream S) = NULL;
-int (*mlx_distributed_all_max_ptr)(mlx_array* res, const mlx_array x, const mlx_distributed_group group , const mlx_stream s) = NULL;
-int (*mlx_distributed_all_min_ptr)(mlx_array* res, const mlx_array x, const mlx_distributed_group group , const mlx_stream s) = NULL;
-int (*mlx_distributed_all_sum_ptr)(mlx_array* res, const mlx_array x, const mlx_distributed_group group , const mlx_stream s) = NULL;
-int (*mlx_distributed_recv_ptr)(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, int src, const mlx_distributed_group group , const mlx_stream s) = NULL;
-int (*mlx_distributed_recv_like_ptr)(mlx_array* res, const mlx_array x, int src, const mlx_distributed_group group , const mlx_stream s) = NULL;
-int (*mlx_distributed_send_ptr)(mlx_array* res, const mlx_array x, int dst, const mlx_distributed_group group , const mlx_stream s) = NULL;
-int (*mlx_distributed_sum_scatter_ptr)(mlx_array* res, const mlx_array x, const mlx_distributed_group group , const mlx_stream s) = NULL;
-int (*mlx_distributed_group_rank_ptr)(mlx_distributed_group group) = NULL;
-int (*mlx_distributed_group_size_ptr)(mlx_distributed_group group) = NULL;
-mlx_distributed_group (*mlx_distributed_group_split_ptr)(mlx_distributed_group group, int color, int key) = NULL;
-bool (*mlx_distributed_is_available_ptr)(void) = NULL;
-mlx_distributed_group (*mlx_distributed_init_ptr)(bool strict) = NULL;
-void (*mlx_set_error_handler_ptr)(mlx_error_handler_func handler, void* data, void (*dtor)(void*)) = NULL;
-void (*_mlx_error_ptr)(const char* file, const int line, const char* fmt, ...) = NULL;
-int (*mlx_export_function_ptr)(const char* file, const mlx_closure fun, const mlx_vector_array args, bool shapeless) = NULL;
-int (*mlx_export_function_kwargs_ptr)(const char* file, const mlx_closure_kwargs fun, const mlx_vector_array args, const mlx_map_string_to_array kwargs, bool shapeless) = NULL;
-mlx_function_exporter (*mlx_function_exporter_new_ptr)(const char* file, const mlx_closure fun, bool shapeless) = NULL;
-int (*mlx_function_exporter_free_ptr)(mlx_function_exporter xfunc) = NULL;
-int (*mlx_function_exporter_apply_ptr)(const mlx_function_exporter xfunc, const mlx_vector_array args) = NULL;
-int (*mlx_function_exporter_apply_kwargs_ptr)(const mlx_function_exporter xfunc, const mlx_vector_array args, const mlx_map_string_to_array kwargs) = NULL;
-mlx_imported_function (*mlx_imported_function_new_ptr)(const char* file) = NULL;
-int (*mlx_imported_function_free_ptr)(mlx_imported_function xfunc) = NULL;
-int (*mlx_imported_function_apply_ptr)(mlx_vector_array* res, const mlx_imported_function xfunc, const mlx_vector_array args) = NULL;
-int (*mlx_imported_function_apply_kwargs_ptr)(mlx_vector_array* res, const mlx_imported_function xfunc, const mlx_vector_array args, const mlx_map_string_to_array kwargs) = NULL;
-mlx_fast_cuda_kernel_config (*mlx_fast_cuda_kernel_config_new_ptr)(void) = NULL;
-void (*mlx_fast_cuda_kernel_config_free_ptr)(mlx_fast_cuda_kernel_config cls) = NULL;
-int (*mlx_fast_cuda_kernel_config_add_output_arg_ptr)(mlx_fast_cuda_kernel_config cls, const int* shape, size_t size, mlx_dtype dtype) = NULL;
-int (*mlx_fast_cuda_kernel_config_set_grid_ptr)(mlx_fast_cuda_kernel_config cls, int grid1, int grid2, int grid3) = NULL;
-int (*mlx_fast_cuda_kernel_config_set_thread_group_ptr)(mlx_fast_cuda_kernel_config cls, int thread1, int thread2, int thread3) = NULL;
-int (*mlx_fast_cuda_kernel_config_set_init_value_ptr)(mlx_fast_cuda_kernel_config cls, float value) = NULL;
-int (*mlx_fast_cuda_kernel_config_set_verbose_ptr)(mlx_fast_cuda_kernel_config cls, bool verbose) = NULL;
-int (*mlx_fast_cuda_kernel_config_add_template_arg_dtype_ptr)(mlx_fast_cuda_kernel_config cls, const char* name, mlx_dtype dtype) = NULL;
-int (*mlx_fast_cuda_kernel_config_add_template_arg_int_ptr)(mlx_fast_cuda_kernel_config cls, const char* name, int value) = NULL;
-int (*mlx_fast_cuda_kernel_config_add_template_arg_bool_ptr)(mlx_fast_cuda_kernel_config cls, const char* name, bool value) = NULL;
-mlx_fast_cuda_kernel (*mlx_fast_cuda_kernel_new_ptr)(const char* name, const mlx_vector_string input_names, const mlx_vector_string output_names, const char* source, const char* header, bool ensure_row_contiguous, int shared_memory) = NULL;
-void (*mlx_fast_cuda_kernel_free_ptr)(mlx_fast_cuda_kernel cls) = NULL;
-int (*mlx_fast_cuda_kernel_apply_ptr)(mlx_vector_array* outputs, mlx_fast_cuda_kernel cls, const mlx_vector_array inputs, const mlx_fast_cuda_kernel_config config, const mlx_stream stream) = NULL;
-int (*mlx_fast_layer_norm_ptr)(mlx_array* res, const mlx_array x, const mlx_array weight , const mlx_array bias , float eps, const mlx_stream s) = NULL;
-mlx_fast_metal_kernel_config (*mlx_fast_metal_kernel_config_new_ptr)(void) = NULL;
-void (*mlx_fast_metal_kernel_config_free_ptr)(mlx_fast_metal_kernel_config cls) = NULL;
-int (*mlx_fast_metal_kernel_config_add_output_arg_ptr)(mlx_fast_metal_kernel_config cls, const int* shape, size_t size, mlx_dtype dtype) = NULL;
-int (*mlx_fast_metal_kernel_config_set_grid_ptr)(mlx_fast_metal_kernel_config cls, int grid1, int grid2, int grid3) = NULL;
-int (*mlx_fast_metal_kernel_config_set_thread_group_ptr)(mlx_fast_metal_kernel_config cls, int thread1, int thread2, int thread3) = NULL;
-int (*mlx_fast_metal_kernel_config_set_init_value_ptr)(mlx_fast_metal_kernel_config cls, float value) = NULL;
-int (*mlx_fast_metal_kernel_config_set_verbose_ptr)(mlx_fast_metal_kernel_config cls, bool verbose) = NULL;
-int (*mlx_fast_metal_kernel_config_add_template_arg_dtype_ptr)(mlx_fast_metal_kernel_config cls, const char* name, mlx_dtype dtype) = NULL;
-int (*mlx_fast_metal_kernel_config_add_template_arg_int_ptr)(mlx_fast_metal_kernel_config cls, const char* name, int value) = NULL;
-int (*mlx_fast_metal_kernel_config_add_template_arg_bool_ptr)(mlx_fast_metal_kernel_config cls, const char* name, bool value) = NULL;
-mlx_fast_metal_kernel (*mlx_fast_metal_kernel_new_ptr)(const char* name, const mlx_vector_string input_names, const mlx_vector_string output_names, const char* source, const char* header, bool ensure_row_contiguous, bool atomic_outputs) = NULL;
-void (*mlx_fast_metal_kernel_free_ptr)(mlx_fast_metal_kernel cls) = NULL;
-int (*mlx_fast_metal_kernel_apply_ptr)(mlx_vector_array* outputs, mlx_fast_metal_kernel cls, const mlx_vector_array inputs, const mlx_fast_metal_kernel_config config, const mlx_stream stream) = NULL;
-int (*mlx_fast_rms_norm_ptr)(mlx_array* res, const mlx_array x, const mlx_array weight , float eps, const mlx_stream s) = NULL;
-int (*mlx_fast_rope_ptr)(mlx_array* res, const mlx_array x, int dims, bool traditional, mlx_optional_float base, float scale, int offset, const mlx_array freqs , const mlx_stream s) = NULL;
-int (*mlx_fast_rope_dynamic_ptr)(mlx_array* res, const mlx_array x, int dims, bool traditional, mlx_optional_float base, float scale, const mlx_array offset, const mlx_array freqs , const mlx_stream s) = NULL;
-int (*mlx_fast_scaled_dot_product_attention_ptr)(mlx_array* res, const mlx_array queries, const mlx_array keys, const mlx_array values, float scale, const char* mask_mode, const mlx_array mask_arr , const mlx_array sinks , const mlx_stream s) = NULL;
-int (*mlx_fft_fft_ptr)(mlx_array* res, const mlx_array a, int n, int axis, const mlx_stream s) = NULL;
-int (*mlx_fft_fft2_ptr)(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s) = NULL;
-int (*mlx_fft_fftn_ptr)(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s) = NULL;
-int (*mlx_fft_fftshift_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, const mlx_stream s) = NULL;
-int (*mlx_fft_ifft_ptr)(mlx_array* res, const mlx_array a, int n, int axis, const mlx_stream s) = NULL;
-int (*mlx_fft_ifft2_ptr)(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s) = NULL;
-int (*mlx_fft_ifftn_ptr)(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s) = NULL;
-int (*mlx_fft_ifftshift_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, const mlx_stream s) = NULL;
-int (*mlx_fft_irfft_ptr)(mlx_array* res, const mlx_array a, int n, int axis, const mlx_stream s) = NULL;
-int (*mlx_fft_irfft2_ptr)(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s) = NULL;
-int (*mlx_fft_irfftn_ptr)(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s) = NULL;
-int (*mlx_fft_rfft_ptr)(mlx_array* res, const mlx_array a, int n, int axis, const mlx_stream s) = NULL;
-int (*mlx_fft_rfft2_ptr)(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s) = NULL;
-int (*mlx_fft_rfftn_ptr)(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s) = NULL;
-int (*mlx_load_reader_ptr)(mlx_array* res, mlx_io_reader in_stream, const mlx_stream s) = NULL;
-int (*mlx_load_ptr)(mlx_array* res, const char* file, const mlx_stream s) = NULL;
-int (*mlx_load_safetensors_reader_ptr)(mlx_map_string_to_array* res_0, mlx_map_string_to_string* res_1, mlx_io_reader in_stream, const mlx_stream s) = NULL;
-int (*mlx_load_safetensors_ptr)(mlx_map_string_to_array* res_0, mlx_map_string_to_string* res_1, const char* file, const mlx_stream s) = NULL;
-int (*mlx_save_writer_ptr)(mlx_io_writer out_stream, const mlx_array a) = NULL;
-int (*mlx_save_ptr)(const char* file, const mlx_array a) = NULL;
-int (*mlx_save_safetensors_writer_ptr)(mlx_io_writer in_stream, const mlx_map_string_to_array param, const mlx_map_string_to_string metadata) = NULL;
-int (*mlx_save_safetensors_ptr)(const char* file, const mlx_map_string_to_array param, const mlx_map_string_to_string metadata) = NULL;
-mlx_io_reader (*mlx_io_reader_new_ptr)(void* desc, mlx_io_vtable vtable) = NULL;
-int (*mlx_io_reader_descriptor_ptr)(void** desc_, mlx_io_reader io) = NULL;
-int (*mlx_io_reader_tostring_ptr)(mlx_string* str_, mlx_io_reader io) = NULL;
-int (*mlx_io_reader_free_ptr)(mlx_io_reader io) = NULL;
-mlx_io_writer (*mlx_io_writer_new_ptr)(void* desc, mlx_io_vtable vtable) = NULL;
-int (*mlx_io_writer_descriptor_ptr)(void** desc_, mlx_io_writer io) = NULL;
-int (*mlx_io_writer_tostring_ptr)(mlx_string* str_, mlx_io_writer io) = NULL;
-int (*mlx_io_writer_free_ptr)(mlx_io_writer io) = NULL;
-int (*mlx_linalg_cholesky_ptr)(mlx_array* res, const mlx_array a, bool upper, const mlx_stream s) = NULL;
-int (*mlx_linalg_cholesky_inv_ptr)(mlx_array* res, const mlx_array a, bool upper, const mlx_stream s) = NULL;
-int (*mlx_linalg_cross_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, int axis, const mlx_stream s) = NULL;
-int (*mlx_linalg_eig_ptr)(mlx_array* res_0, mlx_array* res_1, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_linalg_eigh_ptr)(mlx_array* res_0, mlx_array* res_1, const mlx_array a, const char* UPLO, const mlx_stream s) = NULL;
-int (*mlx_linalg_eigvals_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_linalg_eigvalsh_ptr)(mlx_array* res, const mlx_array a, const char* UPLO, const mlx_stream s) = NULL;
-int (*mlx_linalg_inv_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_linalg_lu_ptr)(mlx_vector_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_linalg_lu_factor_ptr)(mlx_array* res_0, mlx_array* res_1, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_linalg_norm_ptr)(mlx_array* res, const mlx_array a, double ord, const int* axis , size_t axis_num, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_linalg_norm_matrix_ptr)(mlx_array* res, const mlx_array a, const char* ord, const int* axis , size_t axis_num, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_linalg_norm_l2_ptr)(mlx_array* res, const mlx_array a, const int* axis , size_t axis_num, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_linalg_pinv_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_linalg_qr_ptr)(mlx_array* res_0, mlx_array* res_1, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_linalg_solve_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_linalg_solve_triangular_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, bool upper, const mlx_stream s) = NULL;
-int (*mlx_linalg_svd_ptr)(mlx_vector_array* res, const mlx_array a, bool compute_uv, const mlx_stream s) = NULL;
-int (*mlx_linalg_tri_inv_ptr)(mlx_array* res, const mlx_array a, bool upper, const mlx_stream s) = NULL;
-mlx_map_string_to_array (*mlx_map_string_to_array_new_ptr)(void) = NULL;
-int (*mlx_map_string_to_array_set_ptr)(mlx_map_string_to_array* map, const mlx_map_string_to_array src) = NULL;
-int (*mlx_map_string_to_array_free_ptr)(mlx_map_string_to_array map) = NULL;
-int (*mlx_map_string_to_array_insert_ptr)(mlx_map_string_to_array map, const char* key, const mlx_array value) = NULL;
-int (*mlx_map_string_to_array_get_ptr)(mlx_array* value, const mlx_map_string_to_array map, const char* key) = NULL;
-mlx_map_string_to_array_iterator (*mlx_map_string_to_array_iterator_new_ptr)(mlx_map_string_to_array map) = NULL;
-int (*mlx_map_string_to_array_iterator_free_ptr)(mlx_map_string_to_array_iterator it) = NULL;
-int (*mlx_map_string_to_array_iterator_next_ptr)(const char** key, mlx_array* value, mlx_map_string_to_array_iterator it) = NULL;
-mlx_map_string_to_string (*mlx_map_string_to_string_new_ptr)(void) = NULL;
-int (*mlx_map_string_to_string_set_ptr)(mlx_map_string_to_string* map, const mlx_map_string_to_string src) = NULL;
-int (*mlx_map_string_to_string_free_ptr)(mlx_map_string_to_string map) = NULL;
-int (*mlx_map_string_to_string_insert_ptr)(mlx_map_string_to_string map, const char* key, const char* value) = NULL;
-int (*mlx_map_string_to_string_get_ptr)(const char** value, const mlx_map_string_to_string map, const char* key) = NULL;
-mlx_map_string_to_string_iterator (*mlx_map_string_to_string_iterator_new_ptr)(mlx_map_string_to_string map) = NULL;
-int (*mlx_map_string_to_string_iterator_free_ptr)(mlx_map_string_to_string_iterator it) = NULL;
-int (*mlx_map_string_to_string_iterator_next_ptr)(const char** key, const char** value, mlx_map_string_to_string_iterator it) = NULL;
-int (*mlx_clear_cache_ptr)(void) = NULL;
-int (*mlx_get_active_memory_ptr)(size_t* res) = NULL;
-int (*mlx_get_cache_memory_ptr)(size_t* res) = NULL;
-int (*mlx_get_memory_limit_ptr)(size_t* res) = NULL;
-int (*mlx_get_peak_memory_ptr)(size_t* res) = NULL;
-int (*mlx_reset_peak_memory_ptr)(void) = NULL;
-int (*mlx_set_cache_limit_ptr)(size_t* res, size_t limit) = NULL;
-int (*mlx_set_memory_limit_ptr)(size_t* res, size_t limit) = NULL;
-int (*mlx_set_wired_limit_ptr)(size_t* res, size_t limit) = NULL;
-mlx_metal_device_info_t (*mlx_metal_device_info_ptr)(void) = NULL;
-int (*mlx_metal_is_available_ptr)(bool* res) = NULL;
-int (*mlx_metal_start_capture_ptr)(const char* path) = NULL;
-int (*mlx_metal_stop_capture_ptr)(void) = NULL;
-int (*mlx_abs_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_add_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_addmm_ptr)(mlx_array* res, const mlx_array c, const mlx_array a, const mlx_array b, float alpha, float beta, const mlx_stream s) = NULL;
-int (*mlx_all_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_all_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_all_ptr)(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_allclose_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, double rtol, double atol, bool equal_nan, const mlx_stream s) = NULL;
-int (*mlx_any_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_any_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_any_ptr)(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_arange_ptr)(mlx_array* res, double start, double stop, double step, mlx_dtype dtype, const mlx_stream s) = NULL;
-int (*mlx_arccos_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_arccosh_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_arcsin_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_arcsinh_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_arctan_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_arctan2_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_arctanh_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_argmax_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_argmax_ptr)(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_argmin_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_argmin_ptr)(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_argpartition_axis_ptr)(mlx_array* res, const mlx_array a, int kth, int axis, const mlx_stream s) = NULL;
-int (*mlx_argpartition_ptr)(mlx_array* res, const mlx_array a, int kth, const mlx_stream s) = NULL;
-int (*mlx_argsort_axis_ptr)(mlx_array* res, const mlx_array a, int axis, const mlx_stream s) = NULL;
-int (*mlx_argsort_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_array_equal_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, bool equal_nan, const mlx_stream s) = NULL;
-int (*mlx_as_strided_ptr)(mlx_array* res, const mlx_array a, const int* shape, size_t shape_num, const int64_t* strides, size_t strides_num, size_t offset, const mlx_stream s) = NULL;
-int (*mlx_astype_ptr)(mlx_array* res, const mlx_array a, mlx_dtype dtype, const mlx_stream s) = NULL;
-int (*mlx_atleast_1d_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_atleast_2d_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_atleast_3d_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_bitwise_and_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_bitwise_invert_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_bitwise_or_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_bitwise_xor_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_block_masked_mm_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, int block_size, const mlx_array mask_out , const mlx_array mask_lhs , const mlx_array mask_rhs , const mlx_stream s) = NULL;
-int (*mlx_broadcast_arrays_ptr)(mlx_vector_array* res, const mlx_vector_array inputs, const mlx_stream s) = NULL;
-int (*mlx_broadcast_to_ptr)(mlx_array* res, const mlx_array a, const int* shape, size_t shape_num, const mlx_stream s) = NULL;
-int (*mlx_ceil_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_clip_ptr)(mlx_array* res, const mlx_array a, const mlx_array a_min , const mlx_array a_max , const mlx_stream s) = NULL;
-int (*mlx_concatenate_axis_ptr)(mlx_array* res, const mlx_vector_array arrays, int axis, const mlx_stream s) = NULL;
-int (*mlx_concatenate_ptr)(mlx_array* res, const mlx_vector_array arrays, const mlx_stream s) = NULL;
-int (*mlx_conjugate_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_contiguous_ptr)(mlx_array* res, const mlx_array a, bool allow_col_major, const mlx_stream s) = NULL;
-int (*mlx_conv1d_ptr)(mlx_array* res, const mlx_array input, const mlx_array weight, int stride, int padding, int dilation, int groups, const mlx_stream s) = NULL;
-int (*mlx_conv2d_ptr)(mlx_array* res, const mlx_array input, const mlx_array weight, int stride_0, int stride_1, int padding_0, int padding_1, int dilation_0, int dilation_1, int groups, const mlx_stream s) = NULL;
-int (*mlx_conv3d_ptr)(mlx_array* res, const mlx_array input, const mlx_array weight, int stride_0, int stride_1, int stride_2, int padding_0, int padding_1, int padding_2, int dilation_0, int dilation_1, int dilation_2, int groups, const mlx_stream s) = NULL;
-int (*mlx_conv_general_ptr)(mlx_array* res, const mlx_array input, const mlx_array weight, const int* stride, size_t stride_num, const int* padding_lo, size_t padding_lo_num, const int* padding_hi, size_t padding_hi_num, const int* kernel_dilation, size_t kernel_dilation_num, const int* input_dilation, size_t input_dilation_num, int groups, bool flip, const mlx_stream s) = NULL;
-int (*mlx_conv_transpose1d_ptr)(mlx_array* res, const mlx_array input, const mlx_array weight, int stride, int padding, int dilation, int output_padding, int groups, const mlx_stream s) = NULL;
-int (*mlx_conv_transpose2d_ptr)(mlx_array* res, const mlx_array input, const mlx_array weight, int stride_0, int stride_1, int padding_0, int padding_1, int dilation_0, int dilation_1, int output_padding_0, int output_padding_1, int groups, const mlx_stream s) = NULL;
-int (*mlx_conv_transpose3d_ptr)(mlx_array* res, const mlx_array input, const mlx_array weight, int stride_0, int stride_1, int stride_2, int padding_0, int padding_1, int padding_2, int dilation_0, int dilation_1, int dilation_2, int output_padding_0, int output_padding_1, int output_padding_2, int groups, const mlx_stream s) = NULL;
-int (*mlx_copy_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_cos_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_cosh_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_cummax_ptr)(mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s) = NULL;
-int (*mlx_cummin_ptr)(mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s) = NULL;
-int (*mlx_cumprod_ptr)(mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s) = NULL;
-int (*mlx_cumsum_ptr)(mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s) = NULL;
-int (*mlx_degrees_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_depends_ptr)(mlx_vector_array* res, const mlx_vector_array inputs, const mlx_vector_array dependencies) = NULL;
-int (*mlx_dequantize_ptr)(mlx_array* res, const mlx_array w, const mlx_array scales, const mlx_array biases , mlx_optional_int group_size, mlx_optional_int bits, const char* mode, mlx_optional_dtype dtype, const mlx_stream s) = NULL;
-int (*mlx_diag_ptr)(mlx_array* res, const mlx_array a, int k, const mlx_stream s) = NULL;
-int (*mlx_diagonal_ptr)(mlx_array* res, const mlx_array a, int offset, int axis1, int axis2, const mlx_stream s) = NULL;
-int (*mlx_divide_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_divmod_ptr)(mlx_vector_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_einsum_ptr)(mlx_array* res, const char* subscripts, const mlx_vector_array operands, const mlx_stream s) = NULL;
-int (*mlx_equal_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_erf_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_erfinv_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_exp_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_expand_dims_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, const mlx_stream s) = NULL;
-int (*mlx_expand_dims_ptr)(mlx_array* res, const mlx_array a, int axis, const mlx_stream s) = NULL;
-int (*mlx_expm1_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_eye_ptr)(mlx_array* res, int n, int m, int k, mlx_dtype dtype, const mlx_stream s) = NULL;
-int (*mlx_flatten_ptr)(mlx_array* res, const mlx_array a, int start_axis, int end_axis, const mlx_stream s) = NULL;
-int (*mlx_floor_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_floor_divide_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_from_fp8_ptr)(mlx_array* res, const mlx_array x, mlx_dtype dtype, const mlx_stream s) = NULL;
-int (*mlx_full_ptr)(mlx_array* res, const int* shape, size_t shape_num, const mlx_array vals, mlx_dtype dtype, const mlx_stream s) = NULL;
-int (*mlx_full_like_ptr)(mlx_array* res, const mlx_array a, const mlx_array vals, mlx_dtype dtype, const mlx_stream s) = NULL;
-int (*mlx_gather_ptr)(mlx_array* res, const mlx_array a, const mlx_vector_array indices, const int* axes, size_t axes_num, const int* slice_sizes, size_t slice_sizes_num, const mlx_stream s) = NULL;
-int (*mlx_gather_single_ptr)(mlx_array* res, const mlx_array a, const mlx_array indices, int axis, const int* slice_sizes, size_t slice_sizes_num, const mlx_stream s) = NULL;
-int (*mlx_gather_mm_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_array lhs_indices , const mlx_array rhs_indices , bool sorted_indices, const mlx_stream s) = NULL;
-int (*mlx_gather_qmm_ptr)(mlx_array* res, const mlx_array x, const mlx_array w, const mlx_array scales, const mlx_array biases , const mlx_array lhs_indices , const mlx_array rhs_indices , bool transpose, mlx_optional_int group_size, mlx_optional_int bits, const char* mode, bool sorted_indices, const mlx_stream s) = NULL;
-int (*mlx_greater_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_greater_equal_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_hadamard_transform_ptr)(mlx_array* res, const mlx_array a, mlx_optional_float scale, const mlx_stream s) = NULL;
-int (*mlx_identity_ptr)(mlx_array* res, int n, mlx_dtype dtype, const mlx_stream s) = NULL;
-int (*mlx_imag_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_inner_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_isclose_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, double rtol, double atol, bool equal_nan, const mlx_stream s) = NULL;
-int (*mlx_isfinite_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_isinf_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_isnan_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_isneginf_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_isposinf_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_kron_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_left_shift_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_less_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_less_equal_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_linspace_ptr)(mlx_array* res, double start, double stop, int num, mlx_dtype dtype, const mlx_stream s) = NULL;
-int (*mlx_log_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_log10_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_log1p_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_log2_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_logaddexp_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_logcumsumexp_ptr)(mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s) = NULL;
-int (*mlx_logical_and_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_logical_not_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_logical_or_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_logsumexp_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_logsumexp_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_logsumexp_ptr)(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_masked_scatter_ptr)(mlx_array* res, const mlx_array a, const mlx_array mask, const mlx_array src, const mlx_stream s) = NULL;
-int (*mlx_matmul_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_max_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_max_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_max_ptr)(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_maximum_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_mean_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_mean_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_mean_ptr)(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_median_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_meshgrid_ptr)(mlx_vector_array* res, const mlx_vector_array arrays, bool sparse, const char* indexing, const mlx_stream s) = NULL;
-int (*mlx_min_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_min_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_min_ptr)(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_minimum_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_moveaxis_ptr)(mlx_array* res, const mlx_array a, int source, int destination, const mlx_stream s) = NULL;
-int (*mlx_multiply_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_nan_to_num_ptr)(mlx_array* res, const mlx_array a, float nan, mlx_optional_float posinf, mlx_optional_float neginf, const mlx_stream s) = NULL;
-int (*mlx_negative_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_not_equal_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_number_of_elements_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool inverted, mlx_dtype dtype, const mlx_stream s) = NULL;
-int (*mlx_ones_ptr)(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_stream s) = NULL;
-int (*mlx_ones_like_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_outer_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_pad_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, const int* low_pad_size, size_t low_pad_size_num, const int* high_pad_size, size_t high_pad_size_num, const mlx_array pad_value, const char* mode, const mlx_stream s) = NULL;
-int (*mlx_pad_symmetric_ptr)(mlx_array* res, const mlx_array a, int pad_width, const mlx_array pad_value, const char* mode, const mlx_stream s) = NULL;
-int (*mlx_partition_axis_ptr)(mlx_array* res, const mlx_array a, int kth, int axis, const mlx_stream s) = NULL;
-int (*mlx_partition_ptr)(mlx_array* res, const mlx_array a, int kth, const mlx_stream s) = NULL;
-int (*mlx_power_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_prod_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_prod_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_prod_ptr)(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_put_along_axis_ptr)(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array values, int axis, const mlx_stream s) = NULL;
-int (*mlx_qqmm_ptr)(mlx_array* res, const mlx_array x, const mlx_array w, const mlx_array w_scales , mlx_optional_int group_size, mlx_optional_int bits, const char* mode, const mlx_stream s) = NULL;
-int (*mlx_quantize_ptr)(mlx_vector_array* res, const mlx_array w, mlx_optional_int group_size, mlx_optional_int bits, const char* mode, const mlx_stream s) = NULL;
-int (*mlx_quantized_matmul_ptr)(mlx_array* res, const mlx_array x, const mlx_array w, const mlx_array scales, const mlx_array biases , bool transpose, mlx_optional_int group_size, mlx_optional_int bits, const char* mode, const mlx_stream s) = NULL;
-int (*mlx_radians_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_real_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_reciprocal_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_remainder_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_repeat_axis_ptr)(mlx_array* res, const mlx_array arr, int repeats, int axis, const mlx_stream s) = NULL;
-int (*mlx_repeat_ptr)(mlx_array* res, const mlx_array arr, int repeats, const mlx_stream s) = NULL;
-int (*mlx_reshape_ptr)(mlx_array* res, const mlx_array a, const int* shape, size_t shape_num, const mlx_stream s) = NULL;
-int (*mlx_right_shift_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_roll_axis_ptr)(mlx_array* res, const mlx_array a, const int* shift, size_t shift_num, int axis, const mlx_stream s) = NULL;
-int (*mlx_roll_axes_ptr)(mlx_array* res, const mlx_array a, const int* shift, size_t shift_num, const int* axes, size_t axes_num, const mlx_stream s) = NULL;
-int (*mlx_roll_ptr)(mlx_array* res, const mlx_array a, const int* shift, size_t shift_num, const mlx_stream s) = NULL;
-int (*mlx_round_ptr)(mlx_array* res, const mlx_array a, int decimals, const mlx_stream s) = NULL;
-int (*mlx_rsqrt_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_scatter_ptr)(mlx_array* res, const mlx_array a, const mlx_vector_array indices, const mlx_array updates, const int* axes, size_t axes_num, const mlx_stream s) = NULL;
-int (*mlx_scatter_single_ptr)(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array updates, int axis, const mlx_stream s) = NULL;
-int (*mlx_scatter_add_ptr)(mlx_array* res, const mlx_array a, const mlx_vector_array indices, const mlx_array updates, const int* axes, size_t axes_num, const mlx_stream s) = NULL;
-int (*mlx_scatter_add_single_ptr)(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array updates, int axis, const mlx_stream s) = NULL;
-int (*mlx_scatter_add_axis_ptr)(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array values, int axis, const mlx_stream s) = NULL;
-int (*mlx_scatter_max_ptr)(mlx_array* res, const mlx_array a, const mlx_vector_array indices, const mlx_array updates, const int* axes, size_t axes_num, const mlx_stream s) = NULL;
-int (*mlx_scatter_max_single_ptr)(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array updates, int axis, const mlx_stream s) = NULL;
-int (*mlx_scatter_min_ptr)(mlx_array* res, const mlx_array a, const mlx_vector_array indices, const mlx_array updates, const int* axes, size_t axes_num, const mlx_stream s) = NULL;
-int (*mlx_scatter_min_single_ptr)(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array updates, int axis, const mlx_stream s) = NULL;
-int (*mlx_scatter_prod_ptr)(mlx_array* res, const mlx_array a, const mlx_vector_array indices, const mlx_array updates, const int* axes, size_t axes_num, const mlx_stream s) = NULL;
-int (*mlx_scatter_prod_single_ptr)(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array updates, int axis, const mlx_stream s) = NULL;
-int (*mlx_segmented_mm_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_array segments, const mlx_stream s) = NULL;
-int (*mlx_sigmoid_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_sign_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_sin_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_sinh_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_slice_ptr)(mlx_array* res, const mlx_array a, const int* start, size_t start_num, const int* stop, size_t stop_num, const int* strides, size_t strides_num, const mlx_stream s) = NULL;
-int (*mlx_slice_dynamic_ptr)(mlx_array* res, const mlx_array a, const mlx_array start, const int* axes, size_t axes_num, const int* slice_size, size_t slice_size_num, const mlx_stream s) = NULL;
-int (*mlx_slice_update_ptr)(mlx_array* res, const mlx_array src, const mlx_array update, const int* start, size_t start_num, const int* stop, size_t stop_num, const int* strides, size_t strides_num, const mlx_stream s) = NULL;
-int (*mlx_slice_update_dynamic_ptr)(mlx_array* res, const mlx_array src, const mlx_array update, const mlx_array start, const int* axes, size_t axes_num, const mlx_stream s) = NULL;
-int (*mlx_softmax_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool precise, const mlx_stream s) = NULL;
-int (*mlx_softmax_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool precise, const mlx_stream s) = NULL;
-int (*mlx_softmax_ptr)(mlx_array* res, const mlx_array a, bool precise, const mlx_stream s) = NULL;
-int (*mlx_sort_axis_ptr)(mlx_array* res, const mlx_array a, int axis, const mlx_stream s) = NULL;
-int (*mlx_sort_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_split_ptr)(mlx_vector_array* res, const mlx_array a, int num_splits, int axis, const mlx_stream s) = NULL;
-int (*mlx_split_sections_ptr)(mlx_vector_array* res, const mlx_array a, const int* indices, size_t indices_num, int axis, const mlx_stream s) = NULL;
-int (*mlx_sqrt_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_square_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_squeeze_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, const mlx_stream s) = NULL;
-int (*mlx_squeeze_axis_ptr)(mlx_array* res, const mlx_array a, int axis, const mlx_stream s) = NULL;
-int (*mlx_squeeze_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_stack_axis_ptr)(mlx_array* res, const mlx_vector_array arrays, int axis, const mlx_stream s) = NULL;
-int (*mlx_stack_ptr)(mlx_array* res, const mlx_vector_array arrays, const mlx_stream s) = NULL;
-int (*mlx_std_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, int ddof, const mlx_stream s) = NULL;
-int (*mlx_std_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool keepdims, int ddof, const mlx_stream s) = NULL;
-int (*mlx_std_ptr)(mlx_array* res, const mlx_array a, bool keepdims, int ddof, const mlx_stream s) = NULL;
-int (*mlx_stop_gradient_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_subtract_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) = NULL;
-int (*mlx_sum_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_sum_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_sum_ptr)(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s) = NULL;
-int (*mlx_swapaxes_ptr)(mlx_array* res, const mlx_array a, int axis1, int axis2, const mlx_stream s) = NULL;
-int (*mlx_take_axis_ptr)(mlx_array* res, const mlx_array a, const mlx_array indices, int axis, const mlx_stream s) = NULL;
-int (*mlx_take_ptr)(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_stream s) = NULL;
-int (*mlx_take_along_axis_ptr)(mlx_array* res, const mlx_array a, const mlx_array indices, int axis, const mlx_stream s) = NULL;
-int (*mlx_tan_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_tanh_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_tensordot_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const int* axes_a, size_t axes_a_num, const int* axes_b, size_t axes_b_num, const mlx_stream s) = NULL;
-int (*mlx_tensordot_axis_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, int axis, const mlx_stream s) = NULL;
-int (*mlx_tile_ptr)(mlx_array* res, const mlx_array arr, const int* reps, size_t reps_num, const mlx_stream s) = NULL;
-int (*mlx_to_fp8_ptr)(mlx_array* res, const mlx_array x, const mlx_stream s) = NULL;
-int (*mlx_topk_axis_ptr)(mlx_array* res, const mlx_array a, int k, int axis, const mlx_stream s) = NULL;
-int (*mlx_topk_ptr)(mlx_array* res, const mlx_array a, int k, const mlx_stream s) = NULL;
-int (*mlx_trace_ptr)(mlx_array* res, const mlx_array a, int offset, int axis1, int axis2, mlx_dtype dtype, const mlx_stream s) = NULL;
-int (*mlx_transpose_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, const mlx_stream s) = NULL;
-int (*mlx_transpose_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_tri_ptr)(mlx_array* res, int n, int m, int k, mlx_dtype type, const mlx_stream s) = NULL;
-int (*mlx_tril_ptr)(mlx_array* res, const mlx_array x, int k, const mlx_stream s) = NULL;
-int (*mlx_triu_ptr)(mlx_array* res, const mlx_array x, int k, const mlx_stream s) = NULL;
-int (*mlx_unflatten_ptr)(mlx_array* res, const mlx_array a, int axis, const int* shape, size_t shape_num, const mlx_stream s) = NULL;
-int (*mlx_var_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, int ddof, const mlx_stream s) = NULL;
-int (*mlx_var_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool keepdims, int ddof, const mlx_stream s) = NULL;
-int (*mlx_var_ptr)(mlx_array* res, const mlx_array a, bool keepdims, int ddof, const mlx_stream s) = NULL;
-int (*mlx_view_ptr)(mlx_array* res, const mlx_array a, mlx_dtype dtype, const mlx_stream s) = NULL;
-int (*mlx_where_ptr)(mlx_array* res, const mlx_array condition, const mlx_array x, const mlx_array y, const mlx_stream s) = NULL;
-int (*mlx_zeros_ptr)(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_stream s) = NULL;
-int (*mlx_zeros_like_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL;
-int (*mlx_random_bernoulli_ptr)(mlx_array* res, const mlx_array p, const int* shape, size_t shape_num, const mlx_array key , const mlx_stream s) = NULL;
-int (*mlx_random_bits_ptr)(mlx_array* res, const int* shape, size_t shape_num, int width, const mlx_array key , const mlx_stream s) = NULL;
-int (*mlx_random_categorical_shape_ptr)(mlx_array* res, const mlx_array logits, int axis, const int* shape, size_t shape_num, const mlx_array key , const mlx_stream s) = NULL;
-int (*mlx_random_categorical_num_samples_ptr)(mlx_array* res, const mlx_array logits_, int axis, int num_samples, const mlx_array key , const mlx_stream s) = NULL;
-int (*mlx_random_categorical_ptr)(mlx_array* res, const mlx_array logits, int axis, const mlx_array key , const mlx_stream s) = NULL;
-int (*mlx_random_gumbel_ptr)(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_array key , const mlx_stream s) = NULL;
-int (*mlx_random_key_ptr)(mlx_array* res, uint64_t seed) = NULL;
-int (*mlx_random_laplace_ptr)(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, float loc, float scale, const mlx_array key , const mlx_stream s) = NULL;
-int (*mlx_random_multivariate_normal_ptr)(mlx_array* res, const mlx_array mean, const mlx_array cov, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_array key , const mlx_stream s) = NULL;
-int (*mlx_random_normal_broadcast_ptr)(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_array loc , const mlx_array scale , const mlx_array key , const mlx_stream s) = NULL;
-int (*mlx_random_normal_ptr)(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, float loc, float scale, const mlx_array key , const mlx_stream s) = NULL;
-int (*mlx_random_permutation_ptr)(mlx_array* res, const mlx_array x, int axis, const mlx_array key , const mlx_stream s) = NULL;
-int (*mlx_random_permutation_arange_ptr)(mlx_array* res, int x, const mlx_array key , const mlx_stream s) = NULL;
-int (*mlx_random_randint_ptr)(mlx_array* res, const mlx_array low, const mlx_array high, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_array key , const mlx_stream s) = NULL;
-int (*mlx_random_seed_ptr)(uint64_t seed) = NULL;
-int (*mlx_random_split_num_ptr)(mlx_array* res, const mlx_array key, int num, const mlx_stream s) = NULL;
-int (*mlx_random_split_ptr)(mlx_array* res_0, mlx_array* res_1, const mlx_array key, const mlx_stream s) = NULL;
-int (*mlx_random_truncated_normal_ptr)(mlx_array* res, const mlx_array lower, const mlx_array upper, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_array key , const mlx_stream s) = NULL;
-int (*mlx_random_uniform_ptr)(mlx_array* res, const mlx_array low, const mlx_array high, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_array key , const mlx_stream s) = NULL;
-mlx_stream (*mlx_stream_new_ptr)(void) = NULL;
-mlx_stream (*mlx_stream_new_device_ptr)(mlx_device dev) = NULL;
-int (*mlx_stream_set_ptr)(mlx_stream* stream, const mlx_stream src) = NULL;
-int (*mlx_stream_free_ptr)(mlx_stream stream) = NULL;
-int (*mlx_stream_tostring_ptr)(mlx_string* str, mlx_stream stream) = NULL;
-bool (*mlx_stream_equal_ptr)(mlx_stream lhs, mlx_stream rhs) = NULL;
-int (*mlx_stream_get_device_ptr)(mlx_device* dev, mlx_stream stream) = NULL;
-int (*mlx_stream_get_index_ptr)(int* index, mlx_stream stream) = NULL;
-int (*mlx_synchronize_ptr)(mlx_stream stream) = NULL;
-int (*mlx_get_default_stream_ptr)(mlx_stream* stream, mlx_device dev) = NULL;
-int (*mlx_set_default_stream_ptr)(mlx_stream stream) = NULL;
-mlx_stream (*mlx_default_cpu_stream_new_ptr)(void) = NULL;
-mlx_stream (*mlx_default_gpu_stream_new_ptr)(void) = NULL;
-mlx_string (*mlx_string_new_ptr)(void) = NULL;
-mlx_string (*mlx_string_new_data_ptr)(const char* str) = NULL;
-int (*mlx_string_set_ptr)(mlx_string* str, const mlx_string src) = NULL;
-const char* (*mlx_string_data_ptr)(mlx_string str) = NULL;
-int (*mlx_string_free_ptr)(mlx_string str) = NULL;
-int (*mlx_async_eval_ptr)(const mlx_vector_array outputs) = NULL;
-int (*mlx_checkpoint_ptr)(mlx_closure* res, const mlx_closure fun) = NULL;
-int (*mlx_custom_function_ptr)(mlx_closure* res, const mlx_closure fun, const mlx_closure_custom fun_vjp , const mlx_closure_custom_jvp fun_jvp , const mlx_closure_custom_vmap fun_vmap) = NULL;
-int (*mlx_custom_vjp_ptr)(mlx_closure* res, const mlx_closure fun, const mlx_closure_custom fun_vjp) = NULL;
-int (*mlx_eval_ptr)(const mlx_vector_array outputs) = NULL;
-int (*mlx_jvp_ptr)(mlx_vector_array* res_0, mlx_vector_array* res_1, const mlx_closure fun, const mlx_vector_array primals, const mlx_vector_array tangents) = NULL;
-int (*mlx_value_and_grad_ptr)(mlx_closure_value_and_grad* res, const mlx_closure fun, const int* argnums, size_t argnums_num) = NULL;
-int (*mlx_vjp_ptr)(mlx_vector_array* res_0, mlx_vector_array* res_1, const mlx_closure fun, const mlx_vector_array primals, const mlx_vector_array cotangents) = NULL;
-int (*mlx_detail_vmap_replace_ptr)(mlx_vector_array* res, const mlx_vector_array inputs, const mlx_vector_array s_inputs, const mlx_vector_array s_outputs, const int* in_axes, size_t in_axes_num, const int* out_axes, size_t out_axes_num) = NULL;
-int (*mlx_detail_vmap_trace_ptr)(mlx_vector_array* res_0, mlx_vector_array* res_1, const mlx_closure fun, const mlx_vector_array inputs, const int* in_axes, size_t in_axes_num) = NULL;
-mlx_vector_array (*mlx_vector_array_new_ptr)(void) = NULL;
-int (*mlx_vector_array_set_ptr)(mlx_vector_array* vec, const mlx_vector_array src) = NULL;
-int (*mlx_vector_array_free_ptr)(mlx_vector_array vec) = NULL;
-mlx_vector_array (*mlx_vector_array_new_data_ptr)(const mlx_array* data, size_t size) = NULL;
-mlx_vector_array (*mlx_vector_array_new_value_ptr)(const mlx_array val) = NULL;
-int (*mlx_vector_array_set_data_ptr)(mlx_vector_array* vec, const mlx_array* data, size_t size) = NULL;
-int (*mlx_vector_array_set_value_ptr)(mlx_vector_array* vec, const mlx_array val) = NULL;
-int (*mlx_vector_array_append_data_ptr)(mlx_vector_array vec, const mlx_array* data, size_t size) = NULL;
-int (*mlx_vector_array_append_value_ptr)(mlx_vector_array vec, const mlx_array val) = NULL;
-size_t (*mlx_vector_array_size_ptr)(mlx_vector_array vec) = NULL;
-int (*mlx_vector_array_get_ptr)(mlx_array* res, const mlx_vector_array vec, size_t idx) = NULL;
-mlx_vector_vector_array (*mlx_vector_vector_array_new_ptr)(void) = NULL;
-int (*mlx_vector_vector_array_set_ptr)(mlx_vector_vector_array* vec, const mlx_vector_vector_array src) = NULL;
-int (*mlx_vector_vector_array_free_ptr)(mlx_vector_vector_array vec) = NULL;
-mlx_vector_vector_array (*mlx_vector_vector_array_new_data_ptr)(const mlx_vector_array* data, size_t size) = NULL;
-mlx_vector_vector_array (*mlx_vector_vector_array_new_value_ptr)(const mlx_vector_array val) = NULL;
-int (*mlx_vector_vector_array_set_data_ptr)(mlx_vector_vector_array* vec, const mlx_vector_array* data, size_t size) = NULL;
-int (*mlx_vector_vector_array_set_value_ptr)(mlx_vector_vector_array* vec, const mlx_vector_array val) = NULL;
-int (*mlx_vector_vector_array_append_data_ptr)(mlx_vector_vector_array vec, const mlx_vector_array* data, size_t size) = NULL;
-int (*mlx_vector_vector_array_append_value_ptr)(mlx_vector_vector_array vec, const mlx_vector_array val) = NULL;
-size_t (*mlx_vector_vector_array_size_ptr)(mlx_vector_vector_array vec) = NULL;
-int (*mlx_vector_vector_array_get_ptr)(mlx_vector_array* res, const mlx_vector_vector_array vec, size_t idx) = NULL;
-mlx_vector_int (*mlx_vector_int_new_ptr)(void) = NULL;
-int (*mlx_vector_int_set_ptr)(mlx_vector_int* vec, const mlx_vector_int src) = NULL;
-int (*mlx_vector_int_free_ptr)(mlx_vector_int vec) = NULL;
-mlx_vector_int (*mlx_vector_int_new_data_ptr)(int* data, size_t size) = NULL;
-mlx_vector_int (*mlx_vector_int_new_value_ptr)(int val) = NULL;
-int (*mlx_vector_int_set_data_ptr)(mlx_vector_int* vec, int* data, size_t size) = NULL;
-int (*mlx_vector_int_set_value_ptr)(mlx_vector_int* vec, int val) = NULL;
-int (*mlx_vector_int_append_data_ptr)(mlx_vector_int vec, int* data, size_t size) = NULL;
-int (*mlx_vector_int_append_value_ptr)(mlx_vector_int vec, int val) = NULL;
-size_t (*mlx_vector_int_size_ptr)(mlx_vector_int vec) = NULL;
-int (*mlx_vector_int_get_ptr)(int* res, const mlx_vector_int vec, size_t idx) = NULL;
-mlx_vector_string (*mlx_vector_string_new_ptr)(void) = NULL;
-int (*mlx_vector_string_set_ptr)(mlx_vector_string* vec, const mlx_vector_string src) = NULL;
-int (*mlx_vector_string_free_ptr)(mlx_vector_string vec) = NULL;
-mlx_vector_string (*mlx_vector_string_new_data_ptr)(const char** data, size_t size) = NULL;
-mlx_vector_string (*mlx_vector_string_new_value_ptr)(const char* val) = NULL;
-int (*mlx_vector_string_set_data_ptr)(mlx_vector_string* vec, const char** data, size_t size) = NULL;
-int (*mlx_vector_string_set_value_ptr)(mlx_vector_string* vec, const char* val) = NULL;
-int (*mlx_vector_string_append_data_ptr)(mlx_vector_string vec, const char** data, size_t size) = NULL;
-int (*mlx_vector_string_append_value_ptr)(mlx_vector_string vec, const char* val) = NULL;
-size_t (*mlx_vector_string_size_ptr)(mlx_vector_string vec) = NULL;
-int (*mlx_vector_string_get_ptr)(char** res, const mlx_vector_string vec, size_t idx) = NULL;
-int (*mlx_version_ptr)(mlx_string* str_) = NULL;
-
-// Initialize all function pointers via dlsym
-int mlx_load_functions(void* handle) {
-    if (handle == NULL) {
-        fprintf(stderr, "MLX: Invalid library handle\n");
-        return -1;
-    }
-
-    mlx_dtype_size_ptr = dlsym(handle, "mlx_dtype_size");
-    if (mlx_dtype_size_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_dtype_size\n");
-        return -1;
-    }
-    mlx_array_tostring_ptr = dlsym(handle, "mlx_array_tostring");
-    if (mlx_array_tostring_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_tostring\n");
-        return -1;
-    }
-    mlx_array_new_ptr = dlsym(handle, "mlx_array_new");
-    if (mlx_array_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_new\n");
-        return -1;
-    }
-    mlx_array_free_ptr = dlsym(handle, "mlx_array_free");
-    if (mlx_array_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_free\n");
-        return -1;
-    }
-    mlx_array_new_bool_ptr = dlsym(handle, "mlx_array_new_bool");
-    if (mlx_array_new_bool_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_new_bool\n");
-        return -1;
-    }
-    mlx_array_new_int_ptr = dlsym(handle, "mlx_array_new_int");
-    if (mlx_array_new_int_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_new_int\n");
-        return -1;
-    }
-    mlx_array_new_float32_ptr = dlsym(handle, "mlx_array_new_float32");
-    if (mlx_array_new_float32_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_new_float32\n");
-        return -1;
-    }
-    mlx_array_new_float_ptr = dlsym(handle, "mlx_array_new_float");
-    if (mlx_array_new_float_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_new_float\n");
-        return -1;
-    }
-    mlx_array_new_float64_ptr = dlsym(handle, "mlx_array_new_float64");
-    if (mlx_array_new_float64_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_new_float64\n");
-        return -1;
-    }
-    mlx_array_new_double_ptr = dlsym(handle, "mlx_array_new_double");
-    if (mlx_array_new_double_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_new_double\n");
-        return -1;
-    }
-    mlx_array_new_complex_ptr = dlsym(handle, "mlx_array_new_complex");
-    if (mlx_array_new_complex_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_new_complex\n");
-        return -1;
-    }
-    mlx_array_new_data_ptr = dlsym(handle, "mlx_array_new_data");
-    if (mlx_array_new_data_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_new_data\n");
-        return -1;
-    }
-    mlx_array_set_ptr = dlsym(handle, "mlx_array_set");
-    if (mlx_array_set_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_set\n");
-        return -1;
-    }
-    mlx_array_set_bool_ptr = dlsym(handle, "mlx_array_set_bool");
-    if (mlx_array_set_bool_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_set_bool\n");
-        return -1;
-    }
-    mlx_array_set_int_ptr = dlsym(handle, "mlx_array_set_int");
-    if (mlx_array_set_int_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_set_int\n");
-        return -1;
-    }
-    mlx_array_set_float32_ptr = dlsym(handle, "mlx_array_set_float32");
-    if (mlx_array_set_float32_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_set_float32\n");
-        return -1;
-    }
-    mlx_array_set_float_ptr = dlsym(handle, "mlx_array_set_float");
-    if (mlx_array_set_float_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_set_float\n");
-        return -1;
-    }
-    mlx_array_set_float64_ptr = dlsym(handle, "mlx_array_set_float64");
-    if (mlx_array_set_float64_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_set_float64\n");
-        return -1;
-    }
-    mlx_array_set_double_ptr = dlsym(handle, "mlx_array_set_double");
-    if (mlx_array_set_double_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_set_double\n");
-        return -1;
-    }
-    mlx_array_set_complex_ptr = dlsym(handle, "mlx_array_set_complex");
-    if (mlx_array_set_complex_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_set_complex\n");
-        return -1;
-    }
-    mlx_array_set_data_ptr = dlsym(handle, "mlx_array_set_data");
-    if (mlx_array_set_data_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_set_data\n");
-        return -1;
-    }
-    mlx_array_itemsize_ptr = dlsym(handle, "mlx_array_itemsize");
-    if (mlx_array_itemsize_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_itemsize\n");
-        return -1;
-    }
-    mlx_array_size_ptr = dlsym(handle, "mlx_array_size");
-    if (mlx_array_size_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_size\n");
-        return -1;
-    }
-    mlx_array_nbytes_ptr = dlsym(handle, "mlx_array_nbytes");
-    if (mlx_array_nbytes_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_nbytes\n");
-        return -1;
-    }
-    mlx_array_ndim_ptr = dlsym(handle, "mlx_array_ndim");
-    if (mlx_array_ndim_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_ndim\n");
-        return -1;
-    }
-    mlx_array_shape_ptr = dlsym(handle, "mlx_array_shape");
-    if (mlx_array_shape_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_shape\n");
-        return -1;
-    }
-    mlx_array_strides_ptr = dlsym(handle, "mlx_array_strides");
-    if (mlx_array_strides_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_strides\n");
-        return -1;
-    }
-    mlx_array_dim_ptr = dlsym(handle, "mlx_array_dim");
-    if (mlx_array_dim_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_dim\n");
-        return -1;
-    }
-    mlx_array_dtype_ptr = dlsym(handle, "mlx_array_dtype");
-    if (mlx_array_dtype_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_dtype\n");
-        return -1;
-    }
-    mlx_array_eval_ptr = dlsym(handle, "mlx_array_eval");
-    if (mlx_array_eval_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_eval\n");
-        return -1;
-    }
-    mlx_array_item_bool_ptr = dlsym(handle, "mlx_array_item_bool");
-    if (mlx_array_item_bool_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_item_bool\n");
-        return -1;
-    }
-    mlx_array_item_uint8_ptr = dlsym(handle, "mlx_array_item_uint8");
-    if (mlx_array_item_uint8_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_item_uint8\n");
-        return -1;
-    }
-    mlx_array_item_uint16_ptr = dlsym(handle, "mlx_array_item_uint16");
-    if (mlx_array_item_uint16_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_item_uint16\n");
-        return -1;
-    }
-    mlx_array_item_uint32_ptr = dlsym(handle, "mlx_array_item_uint32");
-    if (mlx_array_item_uint32_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_item_uint32\n");
-        return -1;
-    }
-    mlx_array_item_uint64_ptr = dlsym(handle, "mlx_array_item_uint64");
-    if (mlx_array_item_uint64_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_item_uint64\n");
-        return -1;
-    }
-    mlx_array_item_int8_ptr = dlsym(handle, "mlx_array_item_int8");
-    if (mlx_array_item_int8_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_item_int8\n");
-        return -1;
-    }
-    mlx_array_item_int16_ptr = dlsym(handle, "mlx_array_item_int16");
-    if (mlx_array_item_int16_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_item_int16\n");
-        return -1;
-    }
-    mlx_array_item_int32_ptr = dlsym(handle, "mlx_array_item_int32");
-    if (mlx_array_item_int32_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_item_int32\n");
-        return -1;
-    }
-    mlx_array_item_int64_ptr = dlsym(handle, "mlx_array_item_int64");
-    if (mlx_array_item_int64_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_item_int64\n");
-        return -1;
-    }
-    mlx_array_item_float32_ptr = dlsym(handle, "mlx_array_item_float32");
-    if (mlx_array_item_float32_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_item_float32\n");
-        return -1;
-    }
-    mlx_array_item_float64_ptr = dlsym(handle, "mlx_array_item_float64");
-    if (mlx_array_item_float64_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_item_float64\n");
-        return -1;
-    }
-    mlx_array_item_complex64_ptr = dlsym(handle, "mlx_array_item_complex64");
-    if (mlx_array_item_complex64_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_item_complex64\n");
-        return -1;
-    }
-#if defined(__aarch64__) || defined(_M_ARM64)
-    mlx_array_item_float16_ptr = dlsym(handle, "mlx_array_item_float16");
-    if (mlx_array_item_float16_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_item_float16\n");
-        return -1;
-    }
-#endif
-#if defined(__aarch64__) || defined(_M_ARM64)
-    mlx_array_item_bfloat16_ptr = dlsym(handle, "mlx_array_item_bfloat16");
-    if (mlx_array_item_bfloat16_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_item_bfloat16\n");
-        return -1;
-    }
-#endif
-    mlx_array_data_bool_ptr = dlsym(handle, "mlx_array_data_bool");
-    if (mlx_array_data_bool_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_data_bool\n");
-        return -1;
-    }
-    mlx_array_data_uint8_ptr = dlsym(handle, "mlx_array_data_uint8");
-    if (mlx_array_data_uint8_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_data_uint8\n");
-        return -1;
-    }
-    mlx_array_data_uint16_ptr = dlsym(handle, "mlx_array_data_uint16");
-    if (mlx_array_data_uint16_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_data_uint16\n");
-        return -1;
-    }
-    mlx_array_data_uint32_ptr = dlsym(handle, "mlx_array_data_uint32");
-    if (mlx_array_data_uint32_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_data_uint32\n");
-        return -1;
-    }
-    mlx_array_data_uint64_ptr = dlsym(handle, "mlx_array_data_uint64");
-    if (mlx_array_data_uint64_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_data_uint64\n");
-        return -1;
-    }
-    mlx_array_data_int8_ptr = dlsym(handle, "mlx_array_data_int8");
-    if (mlx_array_data_int8_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_data_int8\n");
-        return -1;
-    }
-    mlx_array_data_int16_ptr = dlsym(handle, "mlx_array_data_int16");
-    if (mlx_array_data_int16_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_data_int16\n");
-        return -1;
-    }
-    mlx_array_data_int32_ptr = dlsym(handle, "mlx_array_data_int32");
-    if (mlx_array_data_int32_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_data_int32\n");
-        return -1;
-    }
-    mlx_array_data_int64_ptr = dlsym(handle, "mlx_array_data_int64");
-    if (mlx_array_data_int64_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_data_int64\n");
-        return -1;
-    }
-    mlx_array_data_float32_ptr = dlsym(handle, "mlx_array_data_float32");
-    if (mlx_array_data_float32_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_data_float32\n");
-        return -1;
-    }
-    mlx_array_data_float64_ptr = dlsym(handle, "mlx_array_data_float64");
-    if (mlx_array_data_float64_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_data_float64\n");
-        return -1;
-    }
-    mlx_array_data_complex64_ptr = dlsym(handle, "mlx_array_data_complex64");
-    if (mlx_array_data_complex64_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_data_complex64\n");
-        return -1;
-    }
-#if defined(__aarch64__) || defined(_M_ARM64)
-    mlx_array_data_float16_ptr = dlsym(handle, "mlx_array_data_float16");
-    if (mlx_array_data_float16_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_data_float16\n");
-        return -1;
-    }
-#endif
-#if defined(__aarch64__) || defined(_M_ARM64)
-    mlx_array_data_bfloat16_ptr = dlsym(handle, "mlx_array_data_bfloat16");
-    if (mlx_array_data_bfloat16_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_data_bfloat16\n");
-        return -1;
-    }
-#endif
-    _mlx_array_is_available_ptr = dlsym(handle, "_mlx_array_is_available");
-    if (_mlx_array_is_available_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: _mlx_array_is_available\n");
-        return -1;
-    }
-    _mlx_array_wait_ptr = dlsym(handle, "_mlx_array_wait");
-    if (_mlx_array_wait_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: _mlx_array_wait\n");
-        return -1;
-    }
-    _mlx_array_is_contiguous_ptr = dlsym(handle, "_mlx_array_is_contiguous");
-    if (_mlx_array_is_contiguous_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: _mlx_array_is_contiguous\n");
-        return -1;
-    }
-    _mlx_array_is_row_contiguous_ptr = dlsym(handle, "_mlx_array_is_row_contiguous");
-    if (_mlx_array_is_row_contiguous_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: _mlx_array_is_row_contiguous\n");
-        return -1;
-    }
-    _mlx_array_is_col_contiguous_ptr = dlsym(handle, "_mlx_array_is_col_contiguous");
-    if (_mlx_array_is_col_contiguous_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: _mlx_array_is_col_contiguous\n");
-        return -1;
-    }
-    mlx_closure_new_ptr = dlsym(handle, "mlx_closure_new");
-    if (mlx_closure_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_new\n");
-        return -1;
-    }
-    mlx_closure_free_ptr = dlsym(handle, "mlx_closure_free");
-    if (mlx_closure_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_free\n");
-        return -1;
-    }
-    mlx_closure_new_func_ptr = dlsym(handle, "mlx_closure_new_func");
-    if (mlx_closure_new_func_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_new_func\n");
-        return -1;
-    }
-    mlx_closure_new_func_payload_ptr = dlsym(handle, "mlx_closure_new_func_payload");
-    if (mlx_closure_new_func_payload_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_new_func_payload\n");
-        return -1;
-    }
-    mlx_closure_set_ptr = dlsym(handle, "mlx_closure_set");
-    if (mlx_closure_set_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_set\n");
-        return -1;
-    }
-    mlx_closure_apply_ptr = dlsym(handle, "mlx_closure_apply");
-    if (mlx_closure_apply_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_apply\n");
-        return -1;
-    }
-    mlx_closure_new_unary_ptr = dlsym(handle, "mlx_closure_new_unary");
-    if (mlx_closure_new_unary_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_new_unary\n");
-        return -1;
-    }
-    mlx_closure_kwargs_new_ptr = dlsym(handle, "mlx_closure_kwargs_new");
-    if (mlx_closure_kwargs_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_kwargs_new\n");
-        return -1;
-    }
-    mlx_closure_kwargs_free_ptr = dlsym(handle, "mlx_closure_kwargs_free");
-    if (mlx_closure_kwargs_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_kwargs_free\n");
-        return -1;
-    }
-    mlx_closure_kwargs_new_func_ptr = dlsym(handle, "mlx_closure_kwargs_new_func");
-    if (mlx_closure_kwargs_new_func_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_kwargs_new_func\n");
-        return -1;
-    }
-    mlx_closure_kwargs_new_func_payload_ptr = dlsym(handle, "mlx_closure_kwargs_new_func_payload");
-    if (mlx_closure_kwargs_new_func_payload_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_kwargs_new_func_payload\n");
-        return -1;
-    }
-    mlx_closure_kwargs_set_ptr = dlsym(handle, "mlx_closure_kwargs_set");
-    if (mlx_closure_kwargs_set_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_kwargs_set\n");
-        return -1;
-    }
-    mlx_closure_kwargs_apply_ptr = dlsym(handle, "mlx_closure_kwargs_apply");
-    if (mlx_closure_kwargs_apply_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_kwargs_apply\n");
-        return -1;
-    }
-    mlx_closure_value_and_grad_new_ptr = dlsym(handle, "mlx_closure_value_and_grad_new");
-    if (mlx_closure_value_and_grad_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_value_and_grad_new\n");
-        return -1;
-    }
-    mlx_closure_value_and_grad_free_ptr = dlsym(handle, "mlx_closure_value_and_grad_free");
-    if (mlx_closure_value_and_grad_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_value_and_grad_free\n");
-        return -1;
-    }
-    mlx_closure_value_and_grad_new_func_ptr = dlsym(handle, "mlx_closure_value_and_grad_new_func");
-    if (mlx_closure_value_and_grad_new_func_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_value_and_grad_new_func\n");
-        return -1;
-    }
-    mlx_closure_value_and_grad_new_func_payload_ptr = dlsym(handle, "mlx_closure_value_and_grad_new_func_payload");
-    if (mlx_closure_value_and_grad_new_func_payload_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_value_and_grad_new_func_payload\n");
-        return -1;
-    }
-    mlx_closure_value_and_grad_set_ptr = dlsym(handle, "mlx_closure_value_and_grad_set");
-    if (mlx_closure_value_and_grad_set_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_value_and_grad_set\n");
-        return -1;
-    }
-    mlx_closure_value_and_grad_apply_ptr = dlsym(handle, "mlx_closure_value_and_grad_apply");
-    if (mlx_closure_value_and_grad_apply_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_value_and_grad_apply\n");
-        return -1;
-    }
-    mlx_closure_custom_new_ptr = dlsym(handle, "mlx_closure_custom_new");
-    if (mlx_closure_custom_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_custom_new\n");
-        return -1;
-    }
-    mlx_closure_custom_free_ptr = dlsym(handle, "mlx_closure_custom_free");
-    if (mlx_closure_custom_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_custom_free\n");
-        return -1;
-    }
-    mlx_closure_custom_new_func_ptr = dlsym(handle, "mlx_closure_custom_new_func");
-    if (mlx_closure_custom_new_func_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_custom_new_func\n");
-        return -1;
-    }
-    mlx_closure_custom_new_func_payload_ptr = dlsym(handle, "mlx_closure_custom_new_func_payload");
-    if (mlx_closure_custom_new_func_payload_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_custom_new_func_payload\n");
-        return -1;
-    }
-    mlx_closure_custom_set_ptr = dlsym(handle, "mlx_closure_custom_set");
-    if (mlx_closure_custom_set_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_custom_set\n");
-        return -1;
-    }
-    mlx_closure_custom_apply_ptr = dlsym(handle, "mlx_closure_custom_apply");
-    if (mlx_closure_custom_apply_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_custom_apply\n");
-        return -1;
-    }
-    mlx_closure_custom_jvp_new_ptr = dlsym(handle, "mlx_closure_custom_jvp_new");
-    if (mlx_closure_custom_jvp_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_custom_jvp_new\n");
-        return -1;
-    }
-    mlx_closure_custom_jvp_free_ptr = dlsym(handle, "mlx_closure_custom_jvp_free");
-    if (mlx_closure_custom_jvp_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_custom_jvp_free\n");
-        return -1;
-    }
-    mlx_closure_custom_jvp_new_func_ptr = dlsym(handle, "mlx_closure_custom_jvp_new_func");
-    if (mlx_closure_custom_jvp_new_func_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_custom_jvp_new_func\n");
-        return -1;
-    }
-    mlx_closure_custom_jvp_new_func_payload_ptr = dlsym(handle, "mlx_closure_custom_jvp_new_func_payload");
-    if (mlx_closure_custom_jvp_new_func_payload_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_custom_jvp_new_func_payload\n");
-        return -1;
-    }
-    mlx_closure_custom_jvp_set_ptr = dlsym(handle, "mlx_closure_custom_jvp_set");
-    if (mlx_closure_custom_jvp_set_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_custom_jvp_set\n");
-        return -1;
-    }
-    mlx_closure_custom_jvp_apply_ptr = dlsym(handle, "mlx_closure_custom_jvp_apply");
-    if (mlx_closure_custom_jvp_apply_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_custom_jvp_apply\n");
-        return -1;
-    }
-    mlx_closure_custom_vmap_new_ptr = dlsym(handle, "mlx_closure_custom_vmap_new");
-    if (mlx_closure_custom_vmap_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_custom_vmap_new\n");
-        return -1;
-    }
-    mlx_closure_custom_vmap_free_ptr = dlsym(handle, "mlx_closure_custom_vmap_free");
-    if (mlx_closure_custom_vmap_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_custom_vmap_free\n");
-        return -1;
-    }
-    mlx_closure_custom_vmap_new_func_ptr = dlsym(handle, "mlx_closure_custom_vmap_new_func");
-    if (mlx_closure_custom_vmap_new_func_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_custom_vmap_new_func\n");
-        return -1;
-    }
-    mlx_closure_custom_vmap_new_func_payload_ptr = dlsym(handle, "mlx_closure_custom_vmap_new_func_payload");
-    if (mlx_closure_custom_vmap_new_func_payload_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_custom_vmap_new_func_payload\n");
-        return -1;
-    }
-    mlx_closure_custom_vmap_set_ptr = dlsym(handle, "mlx_closure_custom_vmap_set");
-    if (mlx_closure_custom_vmap_set_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_custom_vmap_set\n");
-        return -1;
-    }
-    mlx_closure_custom_vmap_apply_ptr = dlsym(handle, "mlx_closure_custom_vmap_apply");
-    if (mlx_closure_custom_vmap_apply_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_closure_custom_vmap_apply\n");
-        return -1;
-    }
-    mlx_compile_ptr = dlsym(handle, "mlx_compile");
-    if (mlx_compile_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_compile\n");
-        return -1;
-    }
-    mlx_detail_compile_ptr = dlsym(handle, "mlx_detail_compile");
-    if (mlx_detail_compile_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_detail_compile\n");
-        return -1;
-    }
-    mlx_detail_compile_clear_cache_ptr = dlsym(handle, "mlx_detail_compile_clear_cache");
-    if (mlx_detail_compile_clear_cache_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_detail_compile_clear_cache\n");
-        return -1;
-    }
-    mlx_detail_compile_erase_ptr = dlsym(handle, "mlx_detail_compile_erase");
-    if (mlx_detail_compile_erase_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_detail_compile_erase\n");
-        return -1;
-    }
-    mlx_disable_compile_ptr = dlsym(handle, "mlx_disable_compile");
-    if (mlx_disable_compile_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_disable_compile\n");
-        return -1;
-    }
-    mlx_enable_compile_ptr = dlsym(handle, "mlx_enable_compile");
-    if (mlx_enable_compile_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_enable_compile\n");
-        return -1;
-    }
-    mlx_set_compile_mode_ptr = dlsym(handle, "mlx_set_compile_mode");
-    if (mlx_set_compile_mode_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_set_compile_mode\n");
-        return -1;
-    }
-    mlx_device_new_ptr = dlsym(handle, "mlx_device_new");
-    if (mlx_device_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_device_new\n");
-        return -1;
-    }
-    mlx_device_new_type_ptr = dlsym(handle, "mlx_device_new_type");
-    if (mlx_device_new_type_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_device_new_type\n");
-        return -1;
-    }
-    mlx_device_free_ptr = dlsym(handle, "mlx_device_free");
-    if (mlx_device_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_device_free\n");
-        return -1;
-    }
-    mlx_device_set_ptr = dlsym(handle, "mlx_device_set");
-    if (mlx_device_set_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_device_set\n");
-        return -1;
-    }
-    mlx_device_tostring_ptr = dlsym(handle, "mlx_device_tostring");
-    if (mlx_device_tostring_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_device_tostring\n");
-        return -1;
-    }
-    mlx_device_equal_ptr = dlsym(handle, "mlx_device_equal");
-    if (mlx_device_equal_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_device_equal\n");
-        return -1;
-    }
-    mlx_device_get_index_ptr = dlsym(handle, "mlx_device_get_index");
-    if (mlx_device_get_index_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_device_get_index\n");
-        return -1;
-    }
-    mlx_device_get_type_ptr = dlsym(handle, "mlx_device_get_type");
-    if (mlx_device_get_type_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_device_get_type\n");
-        return -1;
-    }
-    mlx_get_default_device_ptr = dlsym(handle, "mlx_get_default_device");
-    if (mlx_get_default_device_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_get_default_device\n");
-        return -1;
-    }
-    mlx_set_default_device_ptr = dlsym(handle, "mlx_set_default_device");
-    if (mlx_set_default_device_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_set_default_device\n");
-        return -1;
-    }
-    mlx_distributed_all_gather_ptr = dlsym(handle, "mlx_distributed_all_gather");
-    if (mlx_distributed_all_gather_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_distributed_all_gather\n");
-        return -1;
-    }
-    mlx_distributed_all_max_ptr = dlsym(handle, "mlx_distributed_all_max");
-    if (mlx_distributed_all_max_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_distributed_all_max\n");
-        return -1;
-    }
-    mlx_distributed_all_min_ptr = dlsym(handle, "mlx_distributed_all_min");
-    if (mlx_distributed_all_min_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_distributed_all_min\n");
-        return -1;
-    }
-    mlx_distributed_all_sum_ptr = dlsym(handle, "mlx_distributed_all_sum");
-    if (mlx_distributed_all_sum_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_distributed_all_sum\n");
-        return -1;
-    }
-    mlx_distributed_recv_ptr = dlsym(handle, "mlx_distributed_recv");
-    if (mlx_distributed_recv_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_distributed_recv\n");
-        return -1;
-    }
-    mlx_distributed_recv_like_ptr = dlsym(handle, "mlx_distributed_recv_like");
-    if (mlx_distributed_recv_like_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_distributed_recv_like\n");
-        return -1;
-    }
-    mlx_distributed_send_ptr = dlsym(handle, "mlx_distributed_send");
-    if (mlx_distributed_send_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_distributed_send\n");
-        return -1;
-    }
-    mlx_distributed_sum_scatter_ptr = dlsym(handle, "mlx_distributed_sum_scatter");
-    if (mlx_distributed_sum_scatter_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_distributed_sum_scatter\n");
-        return -1;
-    }
-    mlx_distributed_group_rank_ptr = dlsym(handle, "mlx_distributed_group_rank");
-    if (mlx_distributed_group_rank_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_distributed_group_rank\n");
-        return -1;
-    }
-    mlx_distributed_group_size_ptr = dlsym(handle, "mlx_distributed_group_size");
-    if (mlx_distributed_group_size_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_distributed_group_size\n");
-        return -1;
-    }
-    mlx_distributed_group_split_ptr = dlsym(handle, "mlx_distributed_group_split");
-    if (mlx_distributed_group_split_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_distributed_group_split\n");
-        return -1;
-    }
-    mlx_distributed_is_available_ptr = dlsym(handle, "mlx_distributed_is_available");
-    if (mlx_distributed_is_available_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_distributed_is_available\n");
-        return -1;
-    }
-    mlx_distributed_init_ptr = dlsym(handle, "mlx_distributed_init");
-    if (mlx_distributed_init_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_distributed_init\n");
-        return -1;
-    }
-    mlx_set_error_handler_ptr = dlsym(handle, "mlx_set_error_handler");
-    if (mlx_set_error_handler_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_set_error_handler\n");
-        return -1;
-    }
-    _mlx_error_ptr = dlsym(handle, "_mlx_error");
-    if (_mlx_error_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: _mlx_error\n");
-        return -1;
-    }
-    mlx_export_function_ptr = dlsym(handle, "mlx_export_function");
-    if (mlx_export_function_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_export_function\n");
-        return -1;
-    }
-    mlx_export_function_kwargs_ptr = dlsym(handle, "mlx_export_function_kwargs");
-    if (mlx_export_function_kwargs_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_export_function_kwargs\n");
-        return -1;
-    }
-    mlx_function_exporter_new_ptr = dlsym(handle, "mlx_function_exporter_new");
-    if (mlx_function_exporter_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_function_exporter_new\n");
-        return -1;
-    }
-    mlx_function_exporter_free_ptr = dlsym(handle, "mlx_function_exporter_free");
-    if (mlx_function_exporter_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_function_exporter_free\n");
-        return -1;
-    }
-    mlx_function_exporter_apply_ptr = dlsym(handle, "mlx_function_exporter_apply");
-    if (mlx_function_exporter_apply_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_function_exporter_apply\n");
-        return -1;
-    }
-    mlx_function_exporter_apply_kwargs_ptr = dlsym(handle, "mlx_function_exporter_apply_kwargs");
-    if (mlx_function_exporter_apply_kwargs_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_function_exporter_apply_kwargs\n");
-        return -1;
-    }
-    mlx_imported_function_new_ptr = dlsym(handle, "mlx_imported_function_new");
-    if (mlx_imported_function_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_imported_function_new\n");
-        return -1;
-    }
-    mlx_imported_function_free_ptr = dlsym(handle, "mlx_imported_function_free");
-    if (mlx_imported_function_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_imported_function_free\n");
-        return -1;
-    }
-    mlx_imported_function_apply_ptr = dlsym(handle, "mlx_imported_function_apply");
-    if (mlx_imported_function_apply_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_imported_function_apply\n");
-        return -1;
-    }
-    mlx_imported_function_apply_kwargs_ptr = dlsym(handle, "mlx_imported_function_apply_kwargs");
-    if (mlx_imported_function_apply_kwargs_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_imported_function_apply_kwargs\n");
-        return -1;
-    }
-    mlx_fast_cuda_kernel_config_new_ptr = dlsym(handle, "mlx_fast_cuda_kernel_config_new");
-    if (mlx_fast_cuda_kernel_config_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_cuda_kernel_config_new\n");
-        return -1;
-    }
-    mlx_fast_cuda_kernel_config_free_ptr = dlsym(handle, "mlx_fast_cuda_kernel_config_free");
-    if (mlx_fast_cuda_kernel_config_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_cuda_kernel_config_free\n");
-        return -1;
-    }
-    mlx_fast_cuda_kernel_config_add_output_arg_ptr = dlsym(handle, "mlx_fast_cuda_kernel_config_add_output_arg");
-    if (mlx_fast_cuda_kernel_config_add_output_arg_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_cuda_kernel_config_add_output_arg\n");
-        return -1;
-    }
-    mlx_fast_cuda_kernel_config_set_grid_ptr = dlsym(handle, "mlx_fast_cuda_kernel_config_set_grid");
-    if (mlx_fast_cuda_kernel_config_set_grid_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_cuda_kernel_config_set_grid\n");
-        return -1;
-    }
-    mlx_fast_cuda_kernel_config_set_thread_group_ptr = dlsym(handle, "mlx_fast_cuda_kernel_config_set_thread_group");
-    if (mlx_fast_cuda_kernel_config_set_thread_group_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_cuda_kernel_config_set_thread_group\n");
-        return -1;
-    }
-    mlx_fast_cuda_kernel_config_set_init_value_ptr = dlsym(handle, "mlx_fast_cuda_kernel_config_set_init_value");
-    if (mlx_fast_cuda_kernel_config_set_init_value_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_cuda_kernel_config_set_init_value\n");
-        return -1;
-    }
-    mlx_fast_cuda_kernel_config_set_verbose_ptr = dlsym(handle, "mlx_fast_cuda_kernel_config_set_verbose");
-    if (mlx_fast_cuda_kernel_config_set_verbose_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_cuda_kernel_config_set_verbose\n");
-        return -1;
-    }
-    mlx_fast_cuda_kernel_config_add_template_arg_dtype_ptr = dlsym(handle, "mlx_fast_cuda_kernel_config_add_template_arg_dtype");
-    if (mlx_fast_cuda_kernel_config_add_template_arg_dtype_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_cuda_kernel_config_add_template_arg_dtype\n");
-        return -1;
-    }
-    mlx_fast_cuda_kernel_config_add_template_arg_int_ptr = dlsym(handle, "mlx_fast_cuda_kernel_config_add_template_arg_int");
-    if (mlx_fast_cuda_kernel_config_add_template_arg_int_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_cuda_kernel_config_add_template_arg_int\n");
-        return -1;
-    }
-    mlx_fast_cuda_kernel_config_add_template_arg_bool_ptr = dlsym(handle, "mlx_fast_cuda_kernel_config_add_template_arg_bool");
-    if (mlx_fast_cuda_kernel_config_add_template_arg_bool_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_cuda_kernel_config_add_template_arg_bool\n");
-        return -1;
-    }
-    mlx_fast_cuda_kernel_new_ptr = dlsym(handle, "mlx_fast_cuda_kernel_new");
-    if (mlx_fast_cuda_kernel_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_cuda_kernel_new\n");
-        return -1;
-    }
-    mlx_fast_cuda_kernel_free_ptr = dlsym(handle, "mlx_fast_cuda_kernel_free");
-    if (mlx_fast_cuda_kernel_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_cuda_kernel_free\n");
-        return -1;
-    }
-    mlx_fast_cuda_kernel_apply_ptr = dlsym(handle, "mlx_fast_cuda_kernel_apply");
-    if (mlx_fast_cuda_kernel_apply_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_cuda_kernel_apply\n");
-        return -1;
-    }
-    mlx_fast_layer_norm_ptr = dlsym(handle, "mlx_fast_layer_norm");
-    if (mlx_fast_layer_norm_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_layer_norm\n");
-        return -1;
-    }
-    mlx_fast_metal_kernel_config_new_ptr = dlsym(handle, "mlx_fast_metal_kernel_config_new");
-    if (mlx_fast_metal_kernel_config_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_metal_kernel_config_new\n");
-        return -1;
-    }
-    mlx_fast_metal_kernel_config_free_ptr = dlsym(handle, "mlx_fast_metal_kernel_config_free");
-    if (mlx_fast_metal_kernel_config_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_metal_kernel_config_free\n");
-        return -1;
-    }
-    mlx_fast_metal_kernel_config_add_output_arg_ptr = dlsym(handle, "mlx_fast_metal_kernel_config_add_output_arg");
-    if (mlx_fast_metal_kernel_config_add_output_arg_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_metal_kernel_config_add_output_arg\n");
-        return -1;
-    }
-    mlx_fast_metal_kernel_config_set_grid_ptr = dlsym(handle, "mlx_fast_metal_kernel_config_set_grid");
-    if (mlx_fast_metal_kernel_config_set_grid_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_metal_kernel_config_set_grid\n");
-        return -1;
-    }
-    mlx_fast_metal_kernel_config_set_thread_group_ptr = dlsym(handle, "mlx_fast_metal_kernel_config_set_thread_group");
-    if (mlx_fast_metal_kernel_config_set_thread_group_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_metal_kernel_config_set_thread_group\n");
-        return -1;
-    }
-    mlx_fast_metal_kernel_config_set_init_value_ptr = dlsym(handle, "mlx_fast_metal_kernel_config_set_init_value");
-    if (mlx_fast_metal_kernel_config_set_init_value_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_metal_kernel_config_set_init_value\n");
-        return -1;
-    }
-    mlx_fast_metal_kernel_config_set_verbose_ptr = dlsym(handle, "mlx_fast_metal_kernel_config_set_verbose");
-    if (mlx_fast_metal_kernel_config_set_verbose_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_metal_kernel_config_set_verbose\n");
-        return -1;
-    }
-    mlx_fast_metal_kernel_config_add_template_arg_dtype_ptr = dlsym(handle, "mlx_fast_metal_kernel_config_add_template_arg_dtype");
-    if (mlx_fast_metal_kernel_config_add_template_arg_dtype_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_metal_kernel_config_add_template_arg_dtype\n");
-        return -1;
-    }
-    mlx_fast_metal_kernel_config_add_template_arg_int_ptr = dlsym(handle, "mlx_fast_metal_kernel_config_add_template_arg_int");
-    if (mlx_fast_metal_kernel_config_add_template_arg_int_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_metal_kernel_config_add_template_arg_int\n");
-        return -1;
-    }
-    mlx_fast_metal_kernel_config_add_template_arg_bool_ptr = dlsym(handle, "mlx_fast_metal_kernel_config_add_template_arg_bool");
-    if (mlx_fast_metal_kernel_config_add_template_arg_bool_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_metal_kernel_config_add_template_arg_bool\n");
-        return -1;
-    }
-    mlx_fast_metal_kernel_new_ptr = dlsym(handle, "mlx_fast_metal_kernel_new");
-    if (mlx_fast_metal_kernel_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_metal_kernel_new\n");
-        return -1;
-    }
-    mlx_fast_metal_kernel_free_ptr = dlsym(handle, "mlx_fast_metal_kernel_free");
-    if (mlx_fast_metal_kernel_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_metal_kernel_free\n");
-        return -1;
-    }
-    mlx_fast_metal_kernel_apply_ptr = dlsym(handle, "mlx_fast_metal_kernel_apply");
-    if (mlx_fast_metal_kernel_apply_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_metal_kernel_apply\n");
-        return -1;
-    }
-    mlx_fast_rms_norm_ptr = dlsym(handle, "mlx_fast_rms_norm");
-    if (mlx_fast_rms_norm_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_rms_norm\n");
-        return -1;
-    }
-    mlx_fast_rope_ptr = dlsym(handle, "mlx_fast_rope");
-    if (mlx_fast_rope_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_rope\n");
-        return -1;
-    }
-    mlx_fast_rope_dynamic_ptr = dlsym(handle, "mlx_fast_rope_dynamic");
-    if (mlx_fast_rope_dynamic_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_rope_dynamic\n");
-        return -1;
-    }
-    mlx_fast_scaled_dot_product_attention_ptr = dlsym(handle, "mlx_fast_scaled_dot_product_attention");
-    if (mlx_fast_scaled_dot_product_attention_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fast_scaled_dot_product_attention\n");
-        return -1;
-    }
-    mlx_fft_fft_ptr = dlsym(handle, "mlx_fft_fft");
-    if (mlx_fft_fft_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fft_fft\n");
-        return -1;
-    }
-    mlx_fft_fft2_ptr = dlsym(handle, "mlx_fft_fft2");
-    if (mlx_fft_fft2_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fft_fft2\n");
-        return -1;
-    }
-    mlx_fft_fftn_ptr = dlsym(handle, "mlx_fft_fftn");
-    if (mlx_fft_fftn_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fft_fftn\n");
-        return -1;
-    }
-    mlx_fft_fftshift_ptr = dlsym(handle, "mlx_fft_fftshift");
-    if (mlx_fft_fftshift_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fft_fftshift\n");
-        return -1;
-    }
-    mlx_fft_ifft_ptr = dlsym(handle, "mlx_fft_ifft");
-    if (mlx_fft_ifft_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fft_ifft\n");
-        return -1;
-    }
-    mlx_fft_ifft2_ptr = dlsym(handle, "mlx_fft_ifft2");
-    if (mlx_fft_ifft2_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fft_ifft2\n");
-        return -1;
-    }
-    mlx_fft_ifftn_ptr = dlsym(handle, "mlx_fft_ifftn");
-    if (mlx_fft_ifftn_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fft_ifftn\n");
-        return -1;
-    }
-    mlx_fft_ifftshift_ptr = dlsym(handle, "mlx_fft_ifftshift");
-    if (mlx_fft_ifftshift_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fft_ifftshift\n");
-        return -1;
-    }
-    mlx_fft_irfft_ptr = dlsym(handle, "mlx_fft_irfft");
-    if (mlx_fft_irfft_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fft_irfft\n");
-        return -1;
-    }
-    mlx_fft_irfft2_ptr = dlsym(handle, "mlx_fft_irfft2");
-    if (mlx_fft_irfft2_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fft_irfft2\n");
-        return -1;
-    }
-    mlx_fft_irfftn_ptr = dlsym(handle, "mlx_fft_irfftn");
-    if (mlx_fft_irfftn_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fft_irfftn\n");
-        return -1;
-    }
-    mlx_fft_rfft_ptr = dlsym(handle, "mlx_fft_rfft");
-    if (mlx_fft_rfft_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fft_rfft\n");
-        return -1;
-    }
-    mlx_fft_rfft2_ptr = dlsym(handle, "mlx_fft_rfft2");
-    if (mlx_fft_rfft2_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fft_rfft2\n");
-        return -1;
-    }
-    mlx_fft_rfftn_ptr = dlsym(handle, "mlx_fft_rfftn");
-    if (mlx_fft_rfftn_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_fft_rfftn\n");
-        return -1;
-    }
-    mlx_load_reader_ptr = dlsym(handle, "mlx_load_reader");
-    if (mlx_load_reader_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_load_reader\n");
-        return -1;
-    }
-    mlx_load_ptr = dlsym(handle, "mlx_load");
-    if (mlx_load_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_load\n");
-        return -1;
-    }
-    mlx_load_safetensors_reader_ptr = dlsym(handle, "mlx_load_safetensors_reader");
-    if (mlx_load_safetensors_reader_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_load_safetensors_reader\n");
-        return -1;
-    }
-    mlx_load_safetensors_ptr = dlsym(handle, "mlx_load_safetensors");
-    if (mlx_load_safetensors_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_load_safetensors\n");
-        return -1;
-    }
-    mlx_save_writer_ptr = dlsym(handle, "mlx_save_writer");
-    if (mlx_save_writer_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_save_writer\n");
-        return -1;
-    }
-    mlx_save_ptr = dlsym(handle, "mlx_save");
-    if (mlx_save_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_save\n");
-        return -1;
-    }
-    mlx_save_safetensors_writer_ptr = dlsym(handle, "mlx_save_safetensors_writer");
-    if (mlx_save_safetensors_writer_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_save_safetensors_writer\n");
-        return -1;
-    }
-    mlx_save_safetensors_ptr = dlsym(handle, "mlx_save_safetensors");
-    if (mlx_save_safetensors_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_save_safetensors\n");
-        return -1;
-    }
-    mlx_io_reader_new_ptr = dlsym(handle, "mlx_io_reader_new");
-    if (mlx_io_reader_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_io_reader_new\n");
-        return -1;
-    }
-    mlx_io_reader_descriptor_ptr = dlsym(handle, "mlx_io_reader_descriptor");
-    if (mlx_io_reader_descriptor_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_io_reader_descriptor\n");
-        return -1;
-    }
-    mlx_io_reader_tostring_ptr = dlsym(handle, "mlx_io_reader_tostring");
-    if (mlx_io_reader_tostring_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_io_reader_tostring\n");
-        return -1;
-    }
-    mlx_io_reader_free_ptr = dlsym(handle, "mlx_io_reader_free");
-    if (mlx_io_reader_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_io_reader_free\n");
-        return -1;
-    }
-    mlx_io_writer_new_ptr = dlsym(handle, "mlx_io_writer_new");
-    if (mlx_io_writer_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_io_writer_new\n");
-        return -1;
-    }
-    mlx_io_writer_descriptor_ptr = dlsym(handle, "mlx_io_writer_descriptor");
-    if (mlx_io_writer_descriptor_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_io_writer_descriptor\n");
-        return -1;
-    }
-    mlx_io_writer_tostring_ptr = dlsym(handle, "mlx_io_writer_tostring");
-    if (mlx_io_writer_tostring_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_io_writer_tostring\n");
-        return -1;
-    }
-    mlx_io_writer_free_ptr = dlsym(handle, "mlx_io_writer_free");
-    if (mlx_io_writer_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_io_writer_free\n");
-        return -1;
-    }
-    mlx_linalg_cholesky_ptr = dlsym(handle, "mlx_linalg_cholesky");
-    if (mlx_linalg_cholesky_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_linalg_cholesky\n");
-        return -1;
-    }
-    mlx_linalg_cholesky_inv_ptr = dlsym(handle, "mlx_linalg_cholesky_inv");
-    if (mlx_linalg_cholesky_inv_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_linalg_cholesky_inv\n");
-        return -1;
-    }
-    mlx_linalg_cross_ptr = dlsym(handle, "mlx_linalg_cross");
-    if (mlx_linalg_cross_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_linalg_cross\n");
-        return -1;
-    }
-    mlx_linalg_eig_ptr = dlsym(handle, "mlx_linalg_eig");
-    if (mlx_linalg_eig_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_linalg_eig\n");
-        return -1;
-    }
-    mlx_linalg_eigh_ptr = dlsym(handle, "mlx_linalg_eigh");
-    if (mlx_linalg_eigh_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_linalg_eigh\n");
-        return -1;
-    }
-    mlx_linalg_eigvals_ptr = dlsym(handle, "mlx_linalg_eigvals");
-    if (mlx_linalg_eigvals_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_linalg_eigvals\n");
-        return -1;
-    }
-    mlx_linalg_eigvalsh_ptr = dlsym(handle, "mlx_linalg_eigvalsh");
-    if (mlx_linalg_eigvalsh_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_linalg_eigvalsh\n");
-        return -1;
-    }
-    mlx_linalg_inv_ptr = dlsym(handle, "mlx_linalg_inv");
-    if (mlx_linalg_inv_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_linalg_inv\n");
-        return -1;
-    }
-    mlx_linalg_lu_ptr = dlsym(handle, "mlx_linalg_lu");
-    if (mlx_linalg_lu_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_linalg_lu\n");
-        return -1;
-    }
-    mlx_linalg_lu_factor_ptr = dlsym(handle, "mlx_linalg_lu_factor");
-    if (mlx_linalg_lu_factor_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_linalg_lu_factor\n");
-        return -1;
-    }
-    mlx_linalg_norm_ptr = dlsym(handle, "mlx_linalg_norm");
-    if (mlx_linalg_norm_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_linalg_norm\n");
-        return -1;
-    }
-    mlx_linalg_norm_matrix_ptr = dlsym(handle, "mlx_linalg_norm_matrix");
-    if (mlx_linalg_norm_matrix_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_linalg_norm_matrix\n");
-        return -1;
-    }
-    mlx_linalg_norm_l2_ptr = dlsym(handle, "mlx_linalg_norm_l2");
-    if (mlx_linalg_norm_l2_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_linalg_norm_l2\n");
-        return -1;
-    }
-    mlx_linalg_pinv_ptr = dlsym(handle, "mlx_linalg_pinv");
-    if (mlx_linalg_pinv_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_linalg_pinv\n");
-        return -1;
-    }
-    mlx_linalg_qr_ptr = dlsym(handle, "mlx_linalg_qr");
-    if (mlx_linalg_qr_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_linalg_qr\n");
-        return -1;
-    }
-    mlx_linalg_solve_ptr = dlsym(handle, "mlx_linalg_solve");
-    if (mlx_linalg_solve_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_linalg_solve\n");
-        return -1;
-    }
-    mlx_linalg_solve_triangular_ptr = dlsym(handle, "mlx_linalg_solve_triangular");
-    if (mlx_linalg_solve_triangular_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_linalg_solve_triangular\n");
-        return -1;
-    }
-    mlx_linalg_svd_ptr = dlsym(handle, "mlx_linalg_svd");
-    if (mlx_linalg_svd_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_linalg_svd\n");
-        return -1;
-    }
-    mlx_linalg_tri_inv_ptr = dlsym(handle, "mlx_linalg_tri_inv");
-    if (mlx_linalg_tri_inv_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_linalg_tri_inv\n");
-        return -1;
-    }
-    mlx_map_string_to_array_new_ptr = dlsym(handle, "mlx_map_string_to_array_new");
-    if (mlx_map_string_to_array_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_map_string_to_array_new\n");
-        return -1;
-    }
-    mlx_map_string_to_array_set_ptr = dlsym(handle, "mlx_map_string_to_array_set");
-    if (mlx_map_string_to_array_set_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_map_string_to_array_set\n");
-        return -1;
-    }
-    mlx_map_string_to_array_free_ptr = dlsym(handle, "mlx_map_string_to_array_free");
-    if (mlx_map_string_to_array_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_map_string_to_array_free\n");
-        return -1;
-    }
-    mlx_map_string_to_array_insert_ptr = dlsym(handle, "mlx_map_string_to_array_insert");
-    if (mlx_map_string_to_array_insert_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_map_string_to_array_insert\n");
-        return -1;
-    }
-    mlx_map_string_to_array_get_ptr = dlsym(handle, "mlx_map_string_to_array_get");
-    if (mlx_map_string_to_array_get_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_map_string_to_array_get\n");
-        return -1;
-    }
-    mlx_map_string_to_array_iterator_new_ptr = dlsym(handle, "mlx_map_string_to_array_iterator_new");
-    if (mlx_map_string_to_array_iterator_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_map_string_to_array_iterator_new\n");
-        return -1;
-    }
-    mlx_map_string_to_array_iterator_free_ptr = dlsym(handle, "mlx_map_string_to_array_iterator_free");
-    if (mlx_map_string_to_array_iterator_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_map_string_to_array_iterator_free\n");
-        return -1;
-    }
-    mlx_map_string_to_array_iterator_next_ptr = dlsym(handle, "mlx_map_string_to_array_iterator_next");
-    if (mlx_map_string_to_array_iterator_next_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_map_string_to_array_iterator_next\n");
-        return -1;
-    }
-    mlx_map_string_to_string_new_ptr = dlsym(handle, "mlx_map_string_to_string_new");
-    if (mlx_map_string_to_string_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_map_string_to_string_new\n");
-        return -1;
-    }
-    mlx_map_string_to_string_set_ptr = dlsym(handle, "mlx_map_string_to_string_set");
-    if (mlx_map_string_to_string_set_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_map_string_to_string_set\n");
-        return -1;
-    }
-    mlx_map_string_to_string_free_ptr = dlsym(handle, "mlx_map_string_to_string_free");
-    if (mlx_map_string_to_string_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_map_string_to_string_free\n");
-        return -1;
-    }
-    mlx_map_string_to_string_insert_ptr = dlsym(handle, "mlx_map_string_to_string_insert");
-    if (mlx_map_string_to_string_insert_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_map_string_to_string_insert\n");
-        return -1;
-    }
-    mlx_map_string_to_string_get_ptr = dlsym(handle, "mlx_map_string_to_string_get");
-    if (mlx_map_string_to_string_get_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_map_string_to_string_get\n");
-        return -1;
-    }
-    mlx_map_string_to_string_iterator_new_ptr = dlsym(handle, "mlx_map_string_to_string_iterator_new");
-    if (mlx_map_string_to_string_iterator_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_map_string_to_string_iterator_new\n");
-        return -1;
-    }
-    mlx_map_string_to_string_iterator_free_ptr = dlsym(handle, "mlx_map_string_to_string_iterator_free");
-    if (mlx_map_string_to_string_iterator_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_map_string_to_string_iterator_free\n");
-        return -1;
-    }
-    mlx_map_string_to_string_iterator_next_ptr = dlsym(handle, "mlx_map_string_to_string_iterator_next");
-    if (mlx_map_string_to_string_iterator_next_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_map_string_to_string_iterator_next\n");
-        return -1;
-    }
-    mlx_clear_cache_ptr = dlsym(handle, "mlx_clear_cache");
-    if (mlx_clear_cache_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_clear_cache\n");
-        return -1;
-    }
-    mlx_get_active_memory_ptr = dlsym(handle, "mlx_get_active_memory");
-    if (mlx_get_active_memory_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_get_active_memory\n");
-        return -1;
-    }
-    mlx_get_cache_memory_ptr = dlsym(handle, "mlx_get_cache_memory");
-    if (mlx_get_cache_memory_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_get_cache_memory\n");
-        return -1;
-    }
-    mlx_get_memory_limit_ptr = dlsym(handle, "mlx_get_memory_limit");
-    if (mlx_get_memory_limit_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_get_memory_limit\n");
-        return -1;
-    }
-    mlx_get_peak_memory_ptr = dlsym(handle, "mlx_get_peak_memory");
-    if (mlx_get_peak_memory_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_get_peak_memory\n");
-        return -1;
-    }
-    mlx_reset_peak_memory_ptr = dlsym(handle, "mlx_reset_peak_memory");
-    if (mlx_reset_peak_memory_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_reset_peak_memory\n");
-        return -1;
-    }
-    mlx_set_cache_limit_ptr = dlsym(handle, "mlx_set_cache_limit");
-    if (mlx_set_cache_limit_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_set_cache_limit\n");
-        return -1;
-    }
-    mlx_set_memory_limit_ptr = dlsym(handle, "mlx_set_memory_limit");
-    if (mlx_set_memory_limit_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_set_memory_limit\n");
-        return -1;
-    }
-    mlx_set_wired_limit_ptr = dlsym(handle, "mlx_set_wired_limit");
-    if (mlx_set_wired_limit_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_set_wired_limit\n");
-        return -1;
-    }
-    mlx_metal_device_info_ptr = dlsym(handle, "mlx_metal_device_info");
-    if (mlx_metal_device_info_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_metal_device_info\n");
-        return -1;
-    }
-    mlx_metal_is_available_ptr = dlsym(handle, "mlx_metal_is_available");
-    if (mlx_metal_is_available_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_metal_is_available\n");
-        return -1;
-    }
-    mlx_metal_start_capture_ptr = dlsym(handle, "mlx_metal_start_capture");
-    if (mlx_metal_start_capture_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_metal_start_capture\n");
-        return -1;
-    }
-    mlx_metal_stop_capture_ptr = dlsym(handle, "mlx_metal_stop_capture");
-    if (mlx_metal_stop_capture_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_metal_stop_capture\n");
-        return -1;
-    }
-    mlx_abs_ptr = dlsym(handle, "mlx_abs");
-    if (mlx_abs_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_abs\n");
-        return -1;
-    }
-    mlx_add_ptr = dlsym(handle, "mlx_add");
-    if (mlx_add_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_add\n");
-        return -1;
-    }
-    mlx_addmm_ptr = dlsym(handle, "mlx_addmm");
-    if (mlx_addmm_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_addmm\n");
-        return -1;
-    }
-    mlx_all_axes_ptr = dlsym(handle, "mlx_all_axes");
-    if (mlx_all_axes_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_all_axes\n");
-        return -1;
-    }
-    mlx_all_axis_ptr = dlsym(handle, "mlx_all_axis");
-    if (mlx_all_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_all_axis\n");
-        return -1;
-    }
-    mlx_all_ptr = dlsym(handle, "mlx_all");
-    if (mlx_all_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_all\n");
-        return -1;
-    }
-    mlx_allclose_ptr = dlsym(handle, "mlx_allclose");
-    if (mlx_allclose_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_allclose\n");
-        return -1;
-    }
-    mlx_any_axes_ptr = dlsym(handle, "mlx_any_axes");
-    if (mlx_any_axes_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_any_axes\n");
-        return -1;
-    }
-    mlx_any_axis_ptr = dlsym(handle, "mlx_any_axis");
-    if (mlx_any_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_any_axis\n");
-        return -1;
-    }
-    mlx_any_ptr = dlsym(handle, "mlx_any");
-    if (mlx_any_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_any\n");
-        return -1;
-    }
-    mlx_arange_ptr = dlsym(handle, "mlx_arange");
-    if (mlx_arange_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_arange\n");
-        return -1;
-    }
-    mlx_arccos_ptr = dlsym(handle, "mlx_arccos");
-    if (mlx_arccos_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_arccos\n");
-        return -1;
-    }
-    mlx_arccosh_ptr = dlsym(handle, "mlx_arccosh");
-    if (mlx_arccosh_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_arccosh\n");
-        return -1;
-    }
-    mlx_arcsin_ptr = dlsym(handle, "mlx_arcsin");
-    if (mlx_arcsin_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_arcsin\n");
-        return -1;
-    }
-    mlx_arcsinh_ptr = dlsym(handle, "mlx_arcsinh");
-    if (mlx_arcsinh_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_arcsinh\n");
-        return -1;
-    }
-    mlx_arctan_ptr = dlsym(handle, "mlx_arctan");
-    if (mlx_arctan_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_arctan\n");
-        return -1;
-    }
-    mlx_arctan2_ptr = dlsym(handle, "mlx_arctan2");
-    if (mlx_arctan2_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_arctan2\n");
-        return -1;
-    }
-    mlx_arctanh_ptr = dlsym(handle, "mlx_arctanh");
-    if (mlx_arctanh_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_arctanh\n");
-        return -1;
-    }
-    mlx_argmax_axis_ptr = dlsym(handle, "mlx_argmax_axis");
-    if (mlx_argmax_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_argmax_axis\n");
-        return -1;
-    }
-    mlx_argmax_ptr = dlsym(handle, "mlx_argmax");
-    if (mlx_argmax_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_argmax\n");
-        return -1;
-    }
-    mlx_argmin_axis_ptr = dlsym(handle, "mlx_argmin_axis");
-    if (mlx_argmin_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_argmin_axis\n");
-        return -1;
-    }
-    mlx_argmin_ptr = dlsym(handle, "mlx_argmin");
-    if (mlx_argmin_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_argmin\n");
-        return -1;
-    }
-    mlx_argpartition_axis_ptr = dlsym(handle, "mlx_argpartition_axis");
-    if (mlx_argpartition_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_argpartition_axis\n");
-        return -1;
-    }
-    mlx_argpartition_ptr = dlsym(handle, "mlx_argpartition");
-    if (mlx_argpartition_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_argpartition\n");
-        return -1;
-    }
-    mlx_argsort_axis_ptr = dlsym(handle, "mlx_argsort_axis");
-    if (mlx_argsort_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_argsort_axis\n");
-        return -1;
-    }
-    mlx_argsort_ptr = dlsym(handle, "mlx_argsort");
-    if (mlx_argsort_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_argsort\n");
-        return -1;
-    }
-    mlx_array_equal_ptr = dlsym(handle, "mlx_array_equal");
-    if (mlx_array_equal_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_array_equal\n");
-        return -1;
-    }
-    mlx_as_strided_ptr = dlsym(handle, "mlx_as_strided");
-    if (mlx_as_strided_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_as_strided\n");
-        return -1;
-    }
-    mlx_astype_ptr = dlsym(handle, "mlx_astype");
-    if (mlx_astype_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_astype\n");
-        return -1;
-    }
-    mlx_atleast_1d_ptr = dlsym(handle, "mlx_atleast_1d");
-    if (mlx_atleast_1d_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_atleast_1d\n");
-        return -1;
-    }
-    mlx_atleast_2d_ptr = dlsym(handle, "mlx_atleast_2d");
-    if (mlx_atleast_2d_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_atleast_2d\n");
-        return -1;
-    }
-    mlx_atleast_3d_ptr = dlsym(handle, "mlx_atleast_3d");
-    if (mlx_atleast_3d_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_atleast_3d\n");
-        return -1;
-    }
-    mlx_bitwise_and_ptr = dlsym(handle, "mlx_bitwise_and");
-    if (mlx_bitwise_and_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_bitwise_and\n");
-        return -1;
-    }
-    mlx_bitwise_invert_ptr = dlsym(handle, "mlx_bitwise_invert");
-    if (mlx_bitwise_invert_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_bitwise_invert\n");
-        return -1;
-    }
-    mlx_bitwise_or_ptr = dlsym(handle, "mlx_bitwise_or");
-    if (mlx_bitwise_or_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_bitwise_or\n");
-        return -1;
-    }
-    mlx_bitwise_xor_ptr = dlsym(handle, "mlx_bitwise_xor");
-    if (mlx_bitwise_xor_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_bitwise_xor\n");
-        return -1;
-    }
-    mlx_block_masked_mm_ptr = dlsym(handle, "mlx_block_masked_mm");
-    if (mlx_block_masked_mm_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_block_masked_mm\n");
-        return -1;
-    }
-    mlx_broadcast_arrays_ptr = dlsym(handle, "mlx_broadcast_arrays");
-    if (mlx_broadcast_arrays_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_broadcast_arrays\n");
-        return -1;
-    }
-    mlx_broadcast_to_ptr = dlsym(handle, "mlx_broadcast_to");
-    if (mlx_broadcast_to_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_broadcast_to\n");
-        return -1;
-    }
-    mlx_ceil_ptr = dlsym(handle, "mlx_ceil");
-    if (mlx_ceil_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_ceil\n");
-        return -1;
-    }
-    mlx_clip_ptr = dlsym(handle, "mlx_clip");
-    if (mlx_clip_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_clip\n");
-        return -1;
-    }
-    mlx_concatenate_axis_ptr = dlsym(handle, "mlx_concatenate_axis");
-    if (mlx_concatenate_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_concatenate_axis\n");
-        return -1;
-    }
-    mlx_concatenate_ptr = dlsym(handle, "mlx_concatenate");
-    if (mlx_concatenate_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_concatenate\n");
-        return -1;
-    }
-    mlx_conjugate_ptr = dlsym(handle, "mlx_conjugate");
-    if (mlx_conjugate_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_conjugate\n");
-        return -1;
-    }
-    mlx_contiguous_ptr = dlsym(handle, "mlx_contiguous");
-    if (mlx_contiguous_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_contiguous\n");
-        return -1;
-    }
-    mlx_conv1d_ptr = dlsym(handle, "mlx_conv1d");
-    if (mlx_conv1d_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_conv1d\n");
-        return -1;
-    }
-    mlx_conv2d_ptr = dlsym(handle, "mlx_conv2d");
-    if (mlx_conv2d_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_conv2d\n");
-        return -1;
-    }
-    mlx_conv3d_ptr = dlsym(handle, "mlx_conv3d");
-    if (mlx_conv3d_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_conv3d\n");
-        return -1;
-    }
-    mlx_conv_general_ptr = dlsym(handle, "mlx_conv_general");
-    if (mlx_conv_general_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_conv_general\n");
-        return -1;
-    }
-    mlx_conv_transpose1d_ptr = dlsym(handle, "mlx_conv_transpose1d");
-    if (mlx_conv_transpose1d_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_conv_transpose1d\n");
-        return -1;
-    }
-    mlx_conv_transpose2d_ptr = dlsym(handle, "mlx_conv_transpose2d");
-    if (mlx_conv_transpose2d_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_conv_transpose2d\n");
-        return -1;
-    }
-    mlx_conv_transpose3d_ptr = dlsym(handle, "mlx_conv_transpose3d");
-    if (mlx_conv_transpose3d_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_conv_transpose3d\n");
-        return -1;
-    }
-    mlx_copy_ptr = dlsym(handle, "mlx_copy");
-    if (mlx_copy_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_copy\n");
-        return -1;
-    }
-    mlx_cos_ptr = dlsym(handle, "mlx_cos");
-    if (mlx_cos_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_cos\n");
-        return -1;
-    }
-    mlx_cosh_ptr = dlsym(handle, "mlx_cosh");
-    if (mlx_cosh_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_cosh\n");
-        return -1;
-    }
-    mlx_cummax_ptr = dlsym(handle, "mlx_cummax");
-    if (mlx_cummax_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_cummax\n");
-        return -1;
-    }
-    mlx_cummin_ptr = dlsym(handle, "mlx_cummin");
-    if (mlx_cummin_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_cummin\n");
-        return -1;
-    }
-    mlx_cumprod_ptr = dlsym(handle, "mlx_cumprod");
-    if (mlx_cumprod_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_cumprod\n");
-        return -1;
-    }
-    mlx_cumsum_ptr = dlsym(handle, "mlx_cumsum");
-    if (mlx_cumsum_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_cumsum\n");
-        return -1;
-    }
-    mlx_degrees_ptr = dlsym(handle, "mlx_degrees");
-    if (mlx_degrees_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_degrees\n");
-        return -1;
-    }
-    mlx_depends_ptr = dlsym(handle, "mlx_depends");
-    if (mlx_depends_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_depends\n");
-        return -1;
-    }
-    mlx_dequantize_ptr = dlsym(handle, "mlx_dequantize");
-    if (mlx_dequantize_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_dequantize\n");
-        return -1;
-    }
-    mlx_diag_ptr = dlsym(handle, "mlx_diag");
-    if (mlx_diag_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_diag\n");
-        return -1;
-    }
-    mlx_diagonal_ptr = dlsym(handle, "mlx_diagonal");
-    if (mlx_diagonal_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_diagonal\n");
-        return -1;
-    }
-    mlx_divide_ptr = dlsym(handle, "mlx_divide");
-    if (mlx_divide_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_divide\n");
-        return -1;
-    }
-    mlx_divmod_ptr = dlsym(handle, "mlx_divmod");
-    if (mlx_divmod_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_divmod\n");
-        return -1;
-    }
-    mlx_einsum_ptr = dlsym(handle, "mlx_einsum");
-    if (mlx_einsum_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_einsum\n");
-        return -1;
-    }
-    mlx_equal_ptr = dlsym(handle, "mlx_equal");
-    if (mlx_equal_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_equal\n");
-        return -1;
-    }
-    mlx_erf_ptr = dlsym(handle, "mlx_erf");
-    if (mlx_erf_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_erf\n");
-        return -1;
-    }
-    mlx_erfinv_ptr = dlsym(handle, "mlx_erfinv");
-    if (mlx_erfinv_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_erfinv\n");
-        return -1;
-    }
-    mlx_exp_ptr = dlsym(handle, "mlx_exp");
-    if (mlx_exp_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_exp\n");
-        return -1;
-    }
-    mlx_expand_dims_axes_ptr = dlsym(handle, "mlx_expand_dims_axes");
-    if (mlx_expand_dims_axes_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_expand_dims_axes\n");
-        return -1;
-    }
-    mlx_expand_dims_ptr = dlsym(handle, "mlx_expand_dims");
-    if (mlx_expand_dims_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_expand_dims\n");
-        return -1;
-    }
-    mlx_expm1_ptr = dlsym(handle, "mlx_expm1");
-    if (mlx_expm1_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_expm1\n");
-        return -1;
-    }
-    mlx_eye_ptr = dlsym(handle, "mlx_eye");
-    if (mlx_eye_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_eye\n");
-        return -1;
-    }
-    mlx_flatten_ptr = dlsym(handle, "mlx_flatten");
-    if (mlx_flatten_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_flatten\n");
-        return -1;
-    }
-    mlx_floor_ptr = dlsym(handle, "mlx_floor");
-    if (mlx_floor_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_floor\n");
-        return -1;
-    }
-    mlx_floor_divide_ptr = dlsym(handle, "mlx_floor_divide");
-    if (mlx_floor_divide_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_floor_divide\n");
-        return -1;
-    }
-    mlx_from_fp8_ptr = dlsym(handle, "mlx_from_fp8");
-    if (mlx_from_fp8_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_from_fp8\n");
-        return -1;
-    }
-    mlx_full_ptr = dlsym(handle, "mlx_full");
-    if (mlx_full_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_full\n");
-        return -1;
-    }
-    mlx_full_like_ptr = dlsym(handle, "mlx_full_like");
-    if (mlx_full_like_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_full_like\n");
-        return -1;
-    }
-    mlx_gather_ptr = dlsym(handle, "mlx_gather");
-    if (mlx_gather_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_gather\n");
-        return -1;
-    }
-    mlx_gather_single_ptr = dlsym(handle, "mlx_gather_single");
-    if (mlx_gather_single_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_gather_single\n");
-        return -1;
-    }
-    mlx_gather_mm_ptr = dlsym(handle, "mlx_gather_mm");
-    if (mlx_gather_mm_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_gather_mm\n");
-        return -1;
-    }
-    mlx_gather_qmm_ptr = dlsym(handle, "mlx_gather_qmm");
-    if (mlx_gather_qmm_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_gather_qmm\n");
-        return -1;
-    }
-    mlx_greater_ptr = dlsym(handle, "mlx_greater");
-    if (mlx_greater_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_greater\n");
-        return -1;
-    }
-    mlx_greater_equal_ptr = dlsym(handle, "mlx_greater_equal");
-    if (mlx_greater_equal_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_greater_equal\n");
-        return -1;
-    }
-    mlx_hadamard_transform_ptr = dlsym(handle, "mlx_hadamard_transform");
-    if (mlx_hadamard_transform_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_hadamard_transform\n");
-        return -1;
-    }
-    mlx_identity_ptr = dlsym(handle, "mlx_identity");
-    if (mlx_identity_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_identity\n");
-        return -1;
-    }
-    mlx_imag_ptr = dlsym(handle, "mlx_imag");
-    if (mlx_imag_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_imag\n");
-        return -1;
-    }
-    mlx_inner_ptr = dlsym(handle, "mlx_inner");
-    if (mlx_inner_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_inner\n");
-        return -1;
-    }
-    mlx_isclose_ptr = dlsym(handle, "mlx_isclose");
-    if (mlx_isclose_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_isclose\n");
-        return -1;
-    }
-    mlx_isfinite_ptr = dlsym(handle, "mlx_isfinite");
-    if (mlx_isfinite_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_isfinite\n");
-        return -1;
-    }
-    mlx_isinf_ptr = dlsym(handle, "mlx_isinf");
-    if (mlx_isinf_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_isinf\n");
-        return -1;
-    }
-    mlx_isnan_ptr = dlsym(handle, "mlx_isnan");
-    if (mlx_isnan_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_isnan\n");
-        return -1;
-    }
-    mlx_isneginf_ptr = dlsym(handle, "mlx_isneginf");
-    if (mlx_isneginf_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_isneginf\n");
-        return -1;
-    }
-    mlx_isposinf_ptr = dlsym(handle, "mlx_isposinf");
-    if (mlx_isposinf_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_isposinf\n");
-        return -1;
-    }
-    mlx_kron_ptr = dlsym(handle, "mlx_kron");
-    if (mlx_kron_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_kron\n");
-        return -1;
-    }
-    mlx_left_shift_ptr = dlsym(handle, "mlx_left_shift");
-    if (mlx_left_shift_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_left_shift\n");
-        return -1;
-    }
-    mlx_less_ptr = dlsym(handle, "mlx_less");
-    if (mlx_less_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_less\n");
-        return -1;
-    }
-    mlx_less_equal_ptr = dlsym(handle, "mlx_less_equal");
-    if (mlx_less_equal_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_less_equal\n");
-        return -1;
-    }
-    mlx_linspace_ptr = dlsym(handle, "mlx_linspace");
-    if (mlx_linspace_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_linspace\n");
-        return -1;
-    }
-    mlx_log_ptr = dlsym(handle, "mlx_log");
-    if (mlx_log_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_log\n");
-        return -1;
-    }
-    mlx_log10_ptr = dlsym(handle, "mlx_log10");
-    if (mlx_log10_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_log10\n");
-        return -1;
-    }
-    mlx_log1p_ptr = dlsym(handle, "mlx_log1p");
-    if (mlx_log1p_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_log1p\n");
-        return -1;
-    }
-    mlx_log2_ptr = dlsym(handle, "mlx_log2");
-    if (mlx_log2_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_log2\n");
-        return -1;
-    }
-    mlx_logaddexp_ptr = dlsym(handle, "mlx_logaddexp");
-    if (mlx_logaddexp_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_logaddexp\n");
-        return -1;
-    }
-    mlx_logcumsumexp_ptr = dlsym(handle, "mlx_logcumsumexp");
-    if (mlx_logcumsumexp_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_logcumsumexp\n");
-        return -1;
-    }
-    mlx_logical_and_ptr = dlsym(handle, "mlx_logical_and");
-    if (mlx_logical_and_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_logical_and\n");
-        return -1;
-    }
-    mlx_logical_not_ptr = dlsym(handle, "mlx_logical_not");
-    if (mlx_logical_not_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_logical_not\n");
-        return -1;
-    }
-    mlx_logical_or_ptr = dlsym(handle, "mlx_logical_or");
-    if (mlx_logical_or_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_logical_or\n");
-        return -1;
-    }
-    mlx_logsumexp_axes_ptr = dlsym(handle, "mlx_logsumexp_axes");
-    if (mlx_logsumexp_axes_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_logsumexp_axes\n");
-        return -1;
-    }
-    mlx_logsumexp_axis_ptr = dlsym(handle, "mlx_logsumexp_axis");
-    if (mlx_logsumexp_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_logsumexp_axis\n");
-        return -1;
-    }
-    mlx_logsumexp_ptr = dlsym(handle, "mlx_logsumexp");
-    if (mlx_logsumexp_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_logsumexp\n");
-        return -1;
-    }
-    mlx_masked_scatter_ptr = dlsym(handle, "mlx_masked_scatter");
-    if (mlx_masked_scatter_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_masked_scatter\n");
-        return -1;
-    }
-    mlx_matmul_ptr = dlsym(handle, "mlx_matmul");
-    if (mlx_matmul_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_matmul\n");
-        return -1;
-    }
-    mlx_max_axes_ptr = dlsym(handle, "mlx_max_axes");
-    if (mlx_max_axes_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_max_axes\n");
-        return -1;
-    }
-    mlx_max_axis_ptr = dlsym(handle, "mlx_max_axis");
-    if (mlx_max_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_max_axis\n");
-        return -1;
-    }
-    mlx_max_ptr = dlsym(handle, "mlx_max");
-    if (mlx_max_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_max\n");
-        return -1;
-    }
-    mlx_maximum_ptr = dlsym(handle, "mlx_maximum");
-    if (mlx_maximum_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_maximum\n");
-        return -1;
-    }
-    mlx_mean_axes_ptr = dlsym(handle, "mlx_mean_axes");
-    if (mlx_mean_axes_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_mean_axes\n");
-        return -1;
-    }
-    mlx_mean_axis_ptr = dlsym(handle, "mlx_mean_axis");
-    if (mlx_mean_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_mean_axis\n");
-        return -1;
-    }
-    mlx_mean_ptr = dlsym(handle, "mlx_mean");
-    if (mlx_mean_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_mean\n");
-        return -1;
-    }
-    mlx_median_ptr = dlsym(handle, "mlx_median");
-    if (mlx_median_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_median\n");
-        return -1;
-    }
-    mlx_meshgrid_ptr = dlsym(handle, "mlx_meshgrid");
-    if (mlx_meshgrid_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_meshgrid\n");
-        return -1;
-    }
-    mlx_min_axes_ptr = dlsym(handle, "mlx_min_axes");
-    if (mlx_min_axes_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_min_axes\n");
-        return -1;
-    }
-    mlx_min_axis_ptr = dlsym(handle, "mlx_min_axis");
-    if (mlx_min_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_min_axis\n");
-        return -1;
-    }
-    mlx_min_ptr = dlsym(handle, "mlx_min");
-    if (mlx_min_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_min\n");
-        return -1;
-    }
-    mlx_minimum_ptr = dlsym(handle, "mlx_minimum");
-    if (mlx_minimum_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_minimum\n");
-        return -1;
-    }
-    mlx_moveaxis_ptr = dlsym(handle, "mlx_moveaxis");
-    if (mlx_moveaxis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_moveaxis\n");
-        return -1;
-    }
-    mlx_multiply_ptr = dlsym(handle, "mlx_multiply");
-    if (mlx_multiply_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_multiply\n");
-        return -1;
-    }
-    mlx_nan_to_num_ptr = dlsym(handle, "mlx_nan_to_num");
-    if (mlx_nan_to_num_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_nan_to_num\n");
-        return -1;
-    }
-    mlx_negative_ptr = dlsym(handle, "mlx_negative");
-    if (mlx_negative_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_negative\n");
-        return -1;
-    }
-    mlx_not_equal_ptr = dlsym(handle, "mlx_not_equal");
-    if (mlx_not_equal_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_not_equal\n");
-        return -1;
-    }
-    mlx_number_of_elements_ptr = dlsym(handle, "mlx_number_of_elements");
-    if (mlx_number_of_elements_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_number_of_elements\n");
-        return -1;
-    }
-    mlx_ones_ptr = dlsym(handle, "mlx_ones");
-    if (mlx_ones_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_ones\n");
-        return -1;
-    }
-    mlx_ones_like_ptr = dlsym(handle, "mlx_ones_like");
-    if (mlx_ones_like_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_ones_like\n");
-        return -1;
-    }
-    mlx_outer_ptr = dlsym(handle, "mlx_outer");
-    if (mlx_outer_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_outer\n");
-        return -1;
-    }
-    mlx_pad_ptr = dlsym(handle, "mlx_pad");
-    if (mlx_pad_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_pad\n");
-        return -1;
-    }
-    mlx_pad_symmetric_ptr = dlsym(handle, "mlx_pad_symmetric");
-    if (mlx_pad_symmetric_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_pad_symmetric\n");
-        return -1;
-    }
-    mlx_partition_axis_ptr = dlsym(handle, "mlx_partition_axis");
-    if (mlx_partition_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_partition_axis\n");
-        return -1;
-    }
-    mlx_partition_ptr = dlsym(handle, "mlx_partition");
-    if (mlx_partition_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_partition\n");
-        return -1;
-    }
-    mlx_power_ptr = dlsym(handle, "mlx_power");
-    if (mlx_power_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_power\n");
-        return -1;
-    }
-    mlx_prod_axes_ptr = dlsym(handle, "mlx_prod_axes");
-    if (mlx_prod_axes_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_prod_axes\n");
-        return -1;
-    }
-    mlx_prod_axis_ptr = dlsym(handle, "mlx_prod_axis");
-    if (mlx_prod_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_prod_axis\n");
-        return -1;
-    }
-    mlx_prod_ptr = dlsym(handle, "mlx_prod");
-    if (mlx_prod_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_prod\n");
-        return -1;
-    }
-    mlx_put_along_axis_ptr = dlsym(handle, "mlx_put_along_axis");
-    if (mlx_put_along_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_put_along_axis\n");
-        return -1;
-    }
-    mlx_qqmm_ptr = dlsym(handle, "mlx_qqmm");
-    if (mlx_qqmm_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_qqmm\n");
-        return -1;
-    }
-    mlx_quantize_ptr = dlsym(handle, "mlx_quantize");
-    if (mlx_quantize_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_quantize\n");
-        return -1;
-    }
-    mlx_quantized_matmul_ptr = dlsym(handle, "mlx_quantized_matmul");
-    if (mlx_quantized_matmul_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_quantized_matmul\n");
-        return -1;
-    }
-    mlx_radians_ptr = dlsym(handle, "mlx_radians");
-    if (mlx_radians_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_radians\n");
-        return -1;
-    }
-    mlx_real_ptr = dlsym(handle, "mlx_real");
-    if (mlx_real_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_real\n");
-        return -1;
-    }
-    mlx_reciprocal_ptr = dlsym(handle, "mlx_reciprocal");
-    if (mlx_reciprocal_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_reciprocal\n");
-        return -1;
-    }
-    mlx_remainder_ptr = dlsym(handle, "mlx_remainder");
-    if (mlx_remainder_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_remainder\n");
-        return -1;
-    }
-    mlx_repeat_axis_ptr = dlsym(handle, "mlx_repeat_axis");
-    if (mlx_repeat_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_repeat_axis\n");
-        return -1;
-    }
-    mlx_repeat_ptr = dlsym(handle, "mlx_repeat");
-    if (mlx_repeat_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_repeat\n");
-        return -1;
-    }
-    mlx_reshape_ptr = dlsym(handle, "mlx_reshape");
-    if (mlx_reshape_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_reshape\n");
-        return -1;
-    }
-    mlx_right_shift_ptr = dlsym(handle, "mlx_right_shift");
-    if (mlx_right_shift_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_right_shift\n");
-        return -1;
-    }
-    mlx_roll_axis_ptr = dlsym(handle, "mlx_roll_axis");
-    if (mlx_roll_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_roll_axis\n");
-        return -1;
-    }
-    mlx_roll_axes_ptr = dlsym(handle, "mlx_roll_axes");
-    if (mlx_roll_axes_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_roll_axes\n");
-        return -1;
-    }
-    mlx_roll_ptr = dlsym(handle, "mlx_roll");
-    if (mlx_roll_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_roll\n");
-        return -1;
-    }
-    mlx_round_ptr = dlsym(handle, "mlx_round");
-    if (mlx_round_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_round\n");
-        return -1;
-    }
-    mlx_rsqrt_ptr = dlsym(handle, "mlx_rsqrt");
-    if (mlx_rsqrt_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_rsqrt\n");
-        return -1;
-    }
-    mlx_scatter_ptr = dlsym(handle, "mlx_scatter");
-    if (mlx_scatter_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_scatter\n");
-        return -1;
-    }
-    mlx_scatter_single_ptr = dlsym(handle, "mlx_scatter_single");
-    if (mlx_scatter_single_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_scatter_single\n");
-        return -1;
-    }
-    mlx_scatter_add_ptr = dlsym(handle, "mlx_scatter_add");
-    if (mlx_scatter_add_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_scatter_add\n");
-        return -1;
-    }
-    mlx_scatter_add_single_ptr = dlsym(handle, "mlx_scatter_add_single");
-    if (mlx_scatter_add_single_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_scatter_add_single\n");
-        return -1;
-    }
-    mlx_scatter_add_axis_ptr = dlsym(handle, "mlx_scatter_add_axis");
-    if (mlx_scatter_add_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_scatter_add_axis\n");
-        return -1;
-    }
-    mlx_scatter_max_ptr = dlsym(handle, "mlx_scatter_max");
-    if (mlx_scatter_max_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_scatter_max\n");
-        return -1;
-    }
-    mlx_scatter_max_single_ptr = dlsym(handle, "mlx_scatter_max_single");
-    if (mlx_scatter_max_single_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_scatter_max_single\n");
-        return -1;
-    }
-    mlx_scatter_min_ptr = dlsym(handle, "mlx_scatter_min");
-    if (mlx_scatter_min_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_scatter_min\n");
-        return -1;
-    }
-    mlx_scatter_min_single_ptr = dlsym(handle, "mlx_scatter_min_single");
-    if (mlx_scatter_min_single_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_scatter_min_single\n");
-        return -1;
-    }
-    mlx_scatter_prod_ptr = dlsym(handle, "mlx_scatter_prod");
-    if (mlx_scatter_prod_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_scatter_prod\n");
-        return -1;
-    }
-    mlx_scatter_prod_single_ptr = dlsym(handle, "mlx_scatter_prod_single");
-    if (mlx_scatter_prod_single_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_scatter_prod_single\n");
-        return -1;
-    }
-    mlx_segmented_mm_ptr = dlsym(handle, "mlx_segmented_mm");
-    if (mlx_segmented_mm_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_segmented_mm\n");
-        return -1;
-    }
-    mlx_sigmoid_ptr = dlsym(handle, "mlx_sigmoid");
-    if (mlx_sigmoid_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_sigmoid\n");
-        return -1;
-    }
-    mlx_sign_ptr = dlsym(handle, "mlx_sign");
-    if (mlx_sign_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_sign\n");
-        return -1;
-    }
-    mlx_sin_ptr = dlsym(handle, "mlx_sin");
-    if (mlx_sin_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_sin\n");
-        return -1;
-    }
-    mlx_sinh_ptr = dlsym(handle, "mlx_sinh");
-    if (mlx_sinh_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_sinh\n");
-        return -1;
-    }
-    mlx_slice_ptr = dlsym(handle, "mlx_slice");
-    if (mlx_slice_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_slice\n");
-        return -1;
-    }
-    mlx_slice_dynamic_ptr = dlsym(handle, "mlx_slice_dynamic");
-    if (mlx_slice_dynamic_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_slice_dynamic\n");
-        return -1;
-    }
-    mlx_slice_update_ptr = dlsym(handle, "mlx_slice_update");
-    if (mlx_slice_update_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_slice_update\n");
-        return -1;
-    }
-    mlx_slice_update_dynamic_ptr = dlsym(handle, "mlx_slice_update_dynamic");
-    if (mlx_slice_update_dynamic_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_slice_update_dynamic\n");
-        return -1;
-    }
-    mlx_softmax_axes_ptr = dlsym(handle, "mlx_softmax_axes");
-    if (mlx_softmax_axes_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_softmax_axes\n");
-        return -1;
-    }
-    mlx_softmax_axis_ptr = dlsym(handle, "mlx_softmax_axis");
-    if (mlx_softmax_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_softmax_axis\n");
-        return -1;
-    }
-    mlx_softmax_ptr = dlsym(handle, "mlx_softmax");
-    if (mlx_softmax_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_softmax\n");
-        return -1;
-    }
-    mlx_sort_axis_ptr = dlsym(handle, "mlx_sort_axis");
-    if (mlx_sort_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_sort_axis\n");
-        return -1;
-    }
-    mlx_sort_ptr = dlsym(handle, "mlx_sort");
-    if (mlx_sort_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_sort\n");
-        return -1;
-    }
-    mlx_split_ptr = dlsym(handle, "mlx_split");
-    if (mlx_split_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_split\n");
-        return -1;
-    }
-    mlx_split_sections_ptr = dlsym(handle, "mlx_split_sections");
-    if (mlx_split_sections_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_split_sections\n");
-        return -1;
-    }
-    mlx_sqrt_ptr = dlsym(handle, "mlx_sqrt");
-    if (mlx_sqrt_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_sqrt\n");
-        return -1;
-    }
-    mlx_square_ptr = dlsym(handle, "mlx_square");
-    if (mlx_square_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_square\n");
-        return -1;
-    }
-    mlx_squeeze_axes_ptr = dlsym(handle, "mlx_squeeze_axes");
-    if (mlx_squeeze_axes_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_squeeze_axes\n");
-        return -1;
-    }
-    mlx_squeeze_axis_ptr = dlsym(handle, "mlx_squeeze_axis");
-    if (mlx_squeeze_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_squeeze_axis\n");
-        return -1;
-    }
-    mlx_squeeze_ptr = dlsym(handle, "mlx_squeeze");
-    if (mlx_squeeze_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_squeeze\n");
-        return -1;
-    }
-    mlx_stack_axis_ptr = dlsym(handle, "mlx_stack_axis");
-    if (mlx_stack_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_stack_axis\n");
-        return -1;
-    }
-    mlx_stack_ptr = dlsym(handle, "mlx_stack");
-    if (mlx_stack_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_stack\n");
-        return -1;
-    }
-    mlx_std_axes_ptr = dlsym(handle, "mlx_std_axes");
-    if (mlx_std_axes_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_std_axes\n");
-        return -1;
-    }
-    mlx_std_axis_ptr = dlsym(handle, "mlx_std_axis");
-    if (mlx_std_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_std_axis\n");
-        return -1;
-    }
-    mlx_std_ptr = dlsym(handle, "mlx_std");
-    if (mlx_std_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_std\n");
-        return -1;
-    }
-    mlx_stop_gradient_ptr = dlsym(handle, "mlx_stop_gradient");
-    if (mlx_stop_gradient_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_stop_gradient\n");
-        return -1;
-    }
-    mlx_subtract_ptr = dlsym(handle, "mlx_subtract");
-    if (mlx_subtract_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_subtract\n");
-        return -1;
-    }
-    mlx_sum_axes_ptr = dlsym(handle, "mlx_sum_axes");
-    if (mlx_sum_axes_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_sum_axes\n");
-        return -1;
-    }
-    mlx_sum_axis_ptr = dlsym(handle, "mlx_sum_axis");
-    if (mlx_sum_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_sum_axis\n");
-        return -1;
-    }
-    mlx_sum_ptr = dlsym(handle, "mlx_sum");
-    if (mlx_sum_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_sum\n");
-        return -1;
-    }
-    mlx_swapaxes_ptr = dlsym(handle, "mlx_swapaxes");
-    if (mlx_swapaxes_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_swapaxes\n");
-        return -1;
-    }
-    mlx_take_axis_ptr = dlsym(handle, "mlx_take_axis");
-    if (mlx_take_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_take_axis\n");
-        return -1;
-    }
-    mlx_take_ptr = dlsym(handle, "mlx_take");
-    if (mlx_take_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_take\n");
-        return -1;
-    }
-    mlx_take_along_axis_ptr = dlsym(handle, "mlx_take_along_axis");
-    if (mlx_take_along_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_take_along_axis\n");
-        return -1;
-    }
-    mlx_tan_ptr = dlsym(handle, "mlx_tan");
-    if (mlx_tan_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_tan\n");
-        return -1;
-    }
-    mlx_tanh_ptr = dlsym(handle, "mlx_tanh");
-    if (mlx_tanh_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_tanh\n");
-        return -1;
-    }
-    mlx_tensordot_ptr = dlsym(handle, "mlx_tensordot");
-    if (mlx_tensordot_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_tensordot\n");
-        return -1;
-    }
-    mlx_tensordot_axis_ptr = dlsym(handle, "mlx_tensordot_axis");
-    if (mlx_tensordot_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_tensordot_axis\n");
-        return -1;
-    }
-    mlx_tile_ptr = dlsym(handle, "mlx_tile");
-    if (mlx_tile_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_tile\n");
-        return -1;
-    }
-    mlx_to_fp8_ptr = dlsym(handle, "mlx_to_fp8");
-    if (mlx_to_fp8_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_to_fp8\n");
-        return -1;
-    }
-    mlx_topk_axis_ptr = dlsym(handle, "mlx_topk_axis");
-    if (mlx_topk_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_topk_axis\n");
-        return -1;
-    }
-    mlx_topk_ptr = dlsym(handle, "mlx_topk");
-    if (mlx_topk_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_topk\n");
-        return -1;
-    }
-    mlx_trace_ptr = dlsym(handle, "mlx_trace");
-    if (mlx_trace_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_trace\n");
-        return -1;
-    }
-    mlx_transpose_axes_ptr = dlsym(handle, "mlx_transpose_axes");
-    if (mlx_transpose_axes_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_transpose_axes\n");
-        return -1;
-    }
-    mlx_transpose_ptr = dlsym(handle, "mlx_transpose");
-    if (mlx_transpose_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_transpose\n");
-        return -1;
-    }
-    mlx_tri_ptr = dlsym(handle, "mlx_tri");
-    if (mlx_tri_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_tri\n");
-        return -1;
-    }
-    mlx_tril_ptr = dlsym(handle, "mlx_tril");
-    if (mlx_tril_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_tril\n");
-        return -1;
-    }
-    mlx_triu_ptr = dlsym(handle, "mlx_triu");
-    if (mlx_triu_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_triu\n");
-        return -1;
-    }
-    mlx_unflatten_ptr = dlsym(handle, "mlx_unflatten");
-    if (mlx_unflatten_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_unflatten\n");
-        return -1;
-    }
-    mlx_var_axes_ptr = dlsym(handle, "mlx_var_axes");
-    if (mlx_var_axes_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_var_axes\n");
-        return -1;
-    }
-    mlx_var_axis_ptr = dlsym(handle, "mlx_var_axis");
-    if (mlx_var_axis_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_var_axis\n");
-        return -1;
-    }
-    mlx_var_ptr = dlsym(handle, "mlx_var");
-    if (mlx_var_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_var\n");
-        return -1;
-    }
-    mlx_view_ptr = dlsym(handle, "mlx_view");
-    if (mlx_view_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_view\n");
-        return -1;
-    }
-    mlx_where_ptr = dlsym(handle, "mlx_where");
-    if (mlx_where_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_where\n");
-        return -1;
-    }
-    mlx_zeros_ptr = dlsym(handle, "mlx_zeros");
-    if (mlx_zeros_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_zeros\n");
-        return -1;
-    }
-    mlx_zeros_like_ptr = dlsym(handle, "mlx_zeros_like");
-    if (mlx_zeros_like_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_zeros_like\n");
-        return -1;
-    }
-    mlx_random_bernoulli_ptr = dlsym(handle, "mlx_random_bernoulli");
-    if (mlx_random_bernoulli_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_random_bernoulli\n");
-        return -1;
-    }
-    mlx_random_bits_ptr = dlsym(handle, "mlx_random_bits");
-    if (mlx_random_bits_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_random_bits\n");
-        return -1;
-    }
-    mlx_random_categorical_shape_ptr = dlsym(handle, "mlx_random_categorical_shape");
-    if (mlx_random_categorical_shape_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_random_categorical_shape\n");
-        return -1;
-    }
-    mlx_random_categorical_num_samples_ptr = dlsym(handle, "mlx_random_categorical_num_samples");
-    if (mlx_random_categorical_num_samples_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_random_categorical_num_samples\n");
-        return -1;
-    }
-    mlx_random_categorical_ptr = dlsym(handle, "mlx_random_categorical");
-    if (mlx_random_categorical_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_random_categorical\n");
-        return -1;
-    }
-    mlx_random_gumbel_ptr = dlsym(handle, "mlx_random_gumbel");
-    if (mlx_random_gumbel_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_random_gumbel\n");
-        return -1;
-    }
-    mlx_random_key_ptr = dlsym(handle, "mlx_random_key");
-    if (mlx_random_key_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_random_key\n");
-        return -1;
-    }
-    mlx_random_laplace_ptr = dlsym(handle, "mlx_random_laplace");
-    if (mlx_random_laplace_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_random_laplace\n");
-        return -1;
-    }
-    mlx_random_multivariate_normal_ptr = dlsym(handle, "mlx_random_multivariate_normal");
-    if (mlx_random_multivariate_normal_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_random_multivariate_normal\n");
-        return -1;
-    }
-    mlx_random_normal_broadcast_ptr = dlsym(handle, "mlx_random_normal_broadcast");
-    if (mlx_random_normal_broadcast_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_random_normal_broadcast\n");
-        return -1;
-    }
-    mlx_random_normal_ptr = dlsym(handle, "mlx_random_normal");
-    if (mlx_random_normal_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_random_normal\n");
-        return -1;
-    }
-    mlx_random_permutation_ptr = dlsym(handle, "mlx_random_permutation");
-    if (mlx_random_permutation_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_random_permutation\n");
-        return -1;
-    }
-    mlx_random_permutation_arange_ptr = dlsym(handle, "mlx_random_permutation_arange");
-    if (mlx_random_permutation_arange_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_random_permutation_arange\n");
-        return -1;
-    }
-    mlx_random_randint_ptr = dlsym(handle, "mlx_random_randint");
-    if (mlx_random_randint_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_random_randint\n");
-        return -1;
-    }
-    mlx_random_seed_ptr = dlsym(handle, "mlx_random_seed");
-    if (mlx_random_seed_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_random_seed\n");
-        return -1;
-    }
-    mlx_random_split_num_ptr = dlsym(handle, "mlx_random_split_num");
-    if (mlx_random_split_num_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_random_split_num\n");
-        return -1;
-    }
-    mlx_random_split_ptr = dlsym(handle, "mlx_random_split");
-    if (mlx_random_split_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_random_split\n");
-        return -1;
-    }
-    mlx_random_truncated_normal_ptr = dlsym(handle, "mlx_random_truncated_normal");
-    if (mlx_random_truncated_normal_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_random_truncated_normal\n");
-        return -1;
-    }
-    mlx_random_uniform_ptr = dlsym(handle, "mlx_random_uniform");
-    if (mlx_random_uniform_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_random_uniform\n");
-        return -1;
-    }
-    mlx_stream_new_ptr = dlsym(handle, "mlx_stream_new");
-    if (mlx_stream_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_stream_new\n");
-        return -1;
-    }
-    mlx_stream_new_device_ptr = dlsym(handle, "mlx_stream_new_device");
-    if (mlx_stream_new_device_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_stream_new_device\n");
-        return -1;
-    }
-    mlx_stream_set_ptr = dlsym(handle, "mlx_stream_set");
-    if (mlx_stream_set_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_stream_set\n");
-        return -1;
-    }
-    mlx_stream_free_ptr = dlsym(handle, "mlx_stream_free");
-    if (mlx_stream_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_stream_free\n");
-        return -1;
-    }
-    mlx_stream_tostring_ptr = dlsym(handle, "mlx_stream_tostring");
-    if (mlx_stream_tostring_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_stream_tostring\n");
-        return -1;
-    }
-    mlx_stream_equal_ptr = dlsym(handle, "mlx_stream_equal");
-    if (mlx_stream_equal_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_stream_equal\n");
-        return -1;
-    }
-    mlx_stream_get_device_ptr = dlsym(handle, "mlx_stream_get_device");
-    if (mlx_stream_get_device_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_stream_get_device\n");
-        return -1;
-    }
-    mlx_stream_get_index_ptr = dlsym(handle, "mlx_stream_get_index");
-    if (mlx_stream_get_index_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_stream_get_index\n");
-        return -1;
-    }
-    mlx_synchronize_ptr = dlsym(handle, "mlx_synchronize");
-    if (mlx_synchronize_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_synchronize\n");
-        return -1;
-    }
-    mlx_get_default_stream_ptr = dlsym(handle, "mlx_get_default_stream");
-    if (mlx_get_default_stream_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_get_default_stream\n");
-        return -1;
-    }
-    mlx_set_default_stream_ptr = dlsym(handle, "mlx_set_default_stream");
-    if (mlx_set_default_stream_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_set_default_stream\n");
-        return -1;
-    }
-    mlx_default_cpu_stream_new_ptr = dlsym(handle, "mlx_default_cpu_stream_new");
-    if (mlx_default_cpu_stream_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_default_cpu_stream_new\n");
-        return -1;
-    }
-    mlx_default_gpu_stream_new_ptr = dlsym(handle, "mlx_default_gpu_stream_new");
-    if (mlx_default_gpu_stream_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_default_gpu_stream_new\n");
-        return -1;
-    }
-    mlx_string_new_ptr = dlsym(handle, "mlx_string_new");
-    if (mlx_string_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_string_new\n");
-        return -1;
-    }
-    mlx_string_new_data_ptr = dlsym(handle, "mlx_string_new_data");
-    if (mlx_string_new_data_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_string_new_data\n");
-        return -1;
-    }
-    mlx_string_set_ptr = dlsym(handle, "mlx_string_set");
-    if (mlx_string_set_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_string_set\n");
-        return -1;
-    }
-    mlx_string_data_ptr = dlsym(handle, "mlx_string_data");
-    if (mlx_string_data_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_string_data\n");
-        return -1;
-    }
-    mlx_string_free_ptr = dlsym(handle, "mlx_string_free");
-    if (mlx_string_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_string_free\n");
-        return -1;
-    }
-    mlx_async_eval_ptr = dlsym(handle, "mlx_async_eval");
-    if (mlx_async_eval_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_async_eval\n");
-        return -1;
-    }
-    mlx_checkpoint_ptr = dlsym(handle, "mlx_checkpoint");
-    if (mlx_checkpoint_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_checkpoint\n");
-        return -1;
-    }
-    mlx_custom_function_ptr = dlsym(handle, "mlx_custom_function");
-    if (mlx_custom_function_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_custom_function\n");
-        return -1;
-    }
-    mlx_custom_vjp_ptr = dlsym(handle, "mlx_custom_vjp");
-    if (mlx_custom_vjp_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_custom_vjp\n");
-        return -1;
-    }
-    mlx_eval_ptr = dlsym(handle, "mlx_eval");
-    if (mlx_eval_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_eval\n");
-        return -1;
-    }
-    mlx_jvp_ptr = dlsym(handle, "mlx_jvp");
-    if (mlx_jvp_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_jvp\n");
-        return -1;
-    }
-    mlx_value_and_grad_ptr = dlsym(handle, "mlx_value_and_grad");
-    if (mlx_value_and_grad_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_value_and_grad\n");
-        return -1;
-    }
-    mlx_vjp_ptr = dlsym(handle, "mlx_vjp");
-    if (mlx_vjp_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vjp\n");
-        return -1;
-    }
-    mlx_detail_vmap_replace_ptr = dlsym(handle, "mlx_detail_vmap_replace");
-    if (mlx_detail_vmap_replace_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_detail_vmap_replace\n");
-        return -1;
-    }
-    mlx_detail_vmap_trace_ptr = dlsym(handle, "mlx_detail_vmap_trace");
-    if (mlx_detail_vmap_trace_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_detail_vmap_trace\n");
-        return -1;
-    }
-    mlx_vector_array_new_ptr = dlsym(handle, "mlx_vector_array_new");
-    if (mlx_vector_array_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_array_new\n");
-        return -1;
-    }
-    mlx_vector_array_set_ptr = dlsym(handle, "mlx_vector_array_set");
-    if (mlx_vector_array_set_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_array_set\n");
-        return -1;
-    }
-    mlx_vector_array_free_ptr = dlsym(handle, "mlx_vector_array_free");
-    if (mlx_vector_array_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_array_free\n");
-        return -1;
-    }
-    mlx_vector_array_new_data_ptr = dlsym(handle, "mlx_vector_array_new_data");
-    if (mlx_vector_array_new_data_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_array_new_data\n");
-        return -1;
-    }
-    mlx_vector_array_new_value_ptr = dlsym(handle, "mlx_vector_array_new_value");
-    if (mlx_vector_array_new_value_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_array_new_value\n");
-        return -1;
-    }
-    mlx_vector_array_set_data_ptr = dlsym(handle, "mlx_vector_array_set_data");
-    if (mlx_vector_array_set_data_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_array_set_data\n");
-        return -1;
-    }
-    mlx_vector_array_set_value_ptr = dlsym(handle, "mlx_vector_array_set_value");
-    if (mlx_vector_array_set_value_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_array_set_value\n");
-        return -1;
-    }
-    mlx_vector_array_append_data_ptr = dlsym(handle, "mlx_vector_array_append_data");
-    if (mlx_vector_array_append_data_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_array_append_data\n");
-        return -1;
-    }
-    mlx_vector_array_append_value_ptr = dlsym(handle, "mlx_vector_array_append_value");
-    if (mlx_vector_array_append_value_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_array_append_value\n");
-        return -1;
-    }
-    mlx_vector_array_size_ptr = dlsym(handle, "mlx_vector_array_size");
-    if (mlx_vector_array_size_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_array_size\n");
-        return -1;
-    }
-    mlx_vector_array_get_ptr = dlsym(handle, "mlx_vector_array_get");
-    if (mlx_vector_array_get_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_array_get\n");
-        return -1;
-    }
-    mlx_vector_vector_array_new_ptr = dlsym(handle, "mlx_vector_vector_array_new");
-    if (mlx_vector_vector_array_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_vector_array_new\n");
-        return -1;
-    }
-    mlx_vector_vector_array_set_ptr = dlsym(handle, "mlx_vector_vector_array_set");
-    if (mlx_vector_vector_array_set_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_vector_array_set\n");
-        return -1;
-    }
-    mlx_vector_vector_array_free_ptr = dlsym(handle, "mlx_vector_vector_array_free");
-    if (mlx_vector_vector_array_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_vector_array_free\n");
-        return -1;
-    }
-    mlx_vector_vector_array_new_data_ptr = dlsym(handle, "mlx_vector_vector_array_new_data");
-    if (mlx_vector_vector_array_new_data_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_vector_array_new_data\n");
-        return -1;
-    }
-    mlx_vector_vector_array_new_value_ptr = dlsym(handle, "mlx_vector_vector_array_new_value");
-    if (mlx_vector_vector_array_new_value_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_vector_array_new_value\n");
-        return -1;
-    }
-    mlx_vector_vector_array_set_data_ptr = dlsym(handle, "mlx_vector_vector_array_set_data");
-    if (mlx_vector_vector_array_set_data_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_vector_array_set_data\n");
-        return -1;
-    }
-    mlx_vector_vector_array_set_value_ptr = dlsym(handle, "mlx_vector_vector_array_set_value");
-    if (mlx_vector_vector_array_set_value_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_vector_array_set_value\n");
-        return -1;
-    }
-    mlx_vector_vector_array_append_data_ptr = dlsym(handle, "mlx_vector_vector_array_append_data");
-    if (mlx_vector_vector_array_append_data_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_vector_array_append_data\n");
-        return -1;
-    }
-    mlx_vector_vector_array_append_value_ptr = dlsym(handle, "mlx_vector_vector_array_append_value");
-    if (mlx_vector_vector_array_append_value_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_vector_array_append_value\n");
-        return -1;
-    }
-    mlx_vector_vector_array_size_ptr = dlsym(handle, "mlx_vector_vector_array_size");
-    if (mlx_vector_vector_array_size_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_vector_array_size\n");
-        return -1;
-    }
-    mlx_vector_vector_array_get_ptr = dlsym(handle, "mlx_vector_vector_array_get");
-    if (mlx_vector_vector_array_get_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_vector_array_get\n");
-        return -1;
-    }
-    mlx_vector_int_new_ptr = dlsym(handle, "mlx_vector_int_new");
-    if (mlx_vector_int_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_int_new\n");
-        return -1;
-    }
-    mlx_vector_int_set_ptr = dlsym(handle, "mlx_vector_int_set");
-    if (mlx_vector_int_set_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_int_set\n");
-        return -1;
-    }
-    mlx_vector_int_free_ptr = dlsym(handle, "mlx_vector_int_free");
-    if (mlx_vector_int_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_int_free\n");
-        return -1;
-    }
-    mlx_vector_int_new_data_ptr = dlsym(handle, "mlx_vector_int_new_data");
-    if (mlx_vector_int_new_data_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_int_new_data\n");
-        return -1;
-    }
-    mlx_vector_int_new_value_ptr = dlsym(handle, "mlx_vector_int_new_value");
-    if (mlx_vector_int_new_value_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_int_new_value\n");
-        return -1;
-    }
-    mlx_vector_int_set_data_ptr = dlsym(handle, "mlx_vector_int_set_data");
-    if (mlx_vector_int_set_data_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_int_set_data\n");
-        return -1;
-    }
-    mlx_vector_int_set_value_ptr = dlsym(handle, "mlx_vector_int_set_value");
-    if (mlx_vector_int_set_value_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_int_set_value\n");
-        return -1;
-    }
-    mlx_vector_int_append_data_ptr = dlsym(handle, "mlx_vector_int_append_data");
-    if (mlx_vector_int_append_data_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_int_append_data\n");
-        return -1;
-    }
-    mlx_vector_int_append_value_ptr = dlsym(handle, "mlx_vector_int_append_value");
-    if (mlx_vector_int_append_value_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_int_append_value\n");
-        return -1;
-    }
-    mlx_vector_int_size_ptr = dlsym(handle, "mlx_vector_int_size");
-    if (mlx_vector_int_size_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_int_size\n");
-        return -1;
-    }
-    mlx_vector_int_get_ptr = dlsym(handle, "mlx_vector_int_get");
-    if (mlx_vector_int_get_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_int_get\n");
-        return -1;
-    }
-    mlx_vector_string_new_ptr = dlsym(handle, "mlx_vector_string_new");
-    if (mlx_vector_string_new_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_string_new\n");
-        return -1;
-    }
-    mlx_vector_string_set_ptr = dlsym(handle, "mlx_vector_string_set");
-    if (mlx_vector_string_set_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_string_set\n");
-        return -1;
-    }
-    mlx_vector_string_free_ptr = dlsym(handle, "mlx_vector_string_free");
-    if (mlx_vector_string_free_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_string_free\n");
-        return -1;
-    }
-    mlx_vector_string_new_data_ptr = dlsym(handle, "mlx_vector_string_new_data");
-    if (mlx_vector_string_new_data_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_string_new_data\n");
-        return -1;
-    }
-    mlx_vector_string_new_value_ptr = dlsym(handle, "mlx_vector_string_new_value");
-    if (mlx_vector_string_new_value_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_string_new_value\n");
-        return -1;
-    }
-    mlx_vector_string_set_data_ptr = dlsym(handle, "mlx_vector_string_set_data");
-    if (mlx_vector_string_set_data_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_string_set_data\n");
-        return -1;
-    }
-    mlx_vector_string_set_value_ptr = dlsym(handle, "mlx_vector_string_set_value");
-    if (mlx_vector_string_set_value_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_string_set_value\n");
-        return -1;
-    }
-    mlx_vector_string_append_data_ptr = dlsym(handle, "mlx_vector_string_append_data");
-    if (mlx_vector_string_append_data_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_string_append_data\n");
-        return -1;
-    }
-    mlx_vector_string_append_value_ptr = dlsym(handle, "mlx_vector_string_append_value");
-    if (mlx_vector_string_append_value_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_string_append_value\n");
-        return -1;
-    }
-    mlx_vector_string_size_ptr = dlsym(handle, "mlx_vector_string_size");
-    if (mlx_vector_string_size_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_string_size\n");
-        return -1;
-    }
-    mlx_vector_string_get_ptr = dlsym(handle, "mlx_vector_string_get");
-    if (mlx_vector_string_get_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_vector_string_get\n");
-        return -1;
-    }
-    mlx_version_ptr = dlsym(handle, "mlx_version");
-    if (mlx_version_ptr == NULL) {
-        fprintf(stderr, "MLX: Failed to load symbol: mlx_version\n");
-        return -1;
-    }
-    return 0;
-}
-
-// Wrapper function implementations that call through function pointers
-size_t mlx_dtype_size(mlx_dtype dtype) {
-    return mlx_dtype_size_ptr(dtype);
-}
-
-int mlx_array_tostring(mlx_string* str, const mlx_array arr) {
-    return mlx_array_tostring_ptr(str, arr);
-}
-
-mlx_array mlx_array_new(void) {
-    return mlx_array_new_ptr();
-}
-
-int mlx_array_free(mlx_array arr) {
-    return mlx_array_free_ptr(arr);
-}
-
-mlx_array mlx_array_new_bool(bool val) {
-    return mlx_array_new_bool_ptr(val);
-}
-
-mlx_array mlx_array_new_int(int val) {
-    return mlx_array_new_int_ptr(val);
-}
-
-mlx_array mlx_array_new_float32(float val) {
-    return mlx_array_new_float32_ptr(val);
-}
-
-mlx_array mlx_array_new_float(float val) {
-    return mlx_array_new_float_ptr(val);
-}
-
-mlx_array mlx_array_new_float64(double val) {
-    return mlx_array_new_float64_ptr(val);
-}
-
-mlx_array mlx_array_new_double(double val) {
-    return mlx_array_new_double_ptr(val);
-}
-
-mlx_array mlx_array_new_complex(float real_val, float imag_val) {
-    return mlx_array_new_complex_ptr(real_val, imag_val);
-}
-
-mlx_array mlx_array_new_data(const void* data, const int* shape, int dim, mlx_dtype dtype) {
-    return mlx_array_new_data_ptr(data, shape, dim, dtype);
-}
-
-int mlx_array_set(mlx_array* arr, const mlx_array src) {
-    return mlx_array_set_ptr(arr, src);
-}
-
-int mlx_array_set_bool(mlx_array* arr, bool val) {
-    return mlx_array_set_bool_ptr(arr, val);
-}
-
-int mlx_array_set_int(mlx_array* arr, int val) {
-    return mlx_array_set_int_ptr(arr, val);
-}
-
-int mlx_array_set_float32(mlx_array* arr, float val) {
-    return mlx_array_set_float32_ptr(arr, val);
-}
-
-int mlx_array_set_float(mlx_array* arr, float val) {
-    return mlx_array_set_float_ptr(arr, val);
-}
-
-int mlx_array_set_float64(mlx_array* arr, double val) {
-    return mlx_array_set_float64_ptr(arr, val);
-}
-
-int mlx_array_set_double(mlx_array* arr, double val) {
-    return mlx_array_set_double_ptr(arr, val);
-}
-
-int mlx_array_set_complex(mlx_array* arr, float real_val, float imag_val) {
-    return mlx_array_set_complex_ptr(arr, real_val, imag_val);
-}
-
-int mlx_array_set_data(mlx_array* arr, const void* data, const int* shape, int dim, mlx_dtype dtype) {
-    return mlx_array_set_data_ptr(arr, data, shape, dim, dtype);
-}
-
-size_t mlx_array_itemsize(const mlx_array arr) {
-    return mlx_array_itemsize_ptr(arr);
-}
-
-size_t mlx_array_size(const mlx_array arr) {
-    return mlx_array_size_ptr(arr);
-}
-
-size_t mlx_array_nbytes(const mlx_array arr) {
-    return mlx_array_nbytes_ptr(arr);
-}
-
-size_t mlx_array_ndim(const mlx_array arr) {
-    return mlx_array_ndim_ptr(arr);
-}
-
-const int* mlx_array_shape(const mlx_array arr) {
-    return mlx_array_shape_ptr(arr);
-}
-
-const size_t* mlx_array_strides(const mlx_array arr) {
-    return mlx_array_strides_ptr(arr);
-}
-
-int mlx_array_dim(const mlx_array arr, int dim) {
-    return mlx_array_dim_ptr(arr, dim);
-}
-
-mlx_dtype mlx_array_dtype(const mlx_array arr) {
-    return mlx_array_dtype_ptr(arr);
-}
-
-int mlx_array_eval(mlx_array arr) {
-    return mlx_array_eval_ptr(arr);
-}
-
-int mlx_array_item_bool(bool* res, const mlx_array arr) {
-    return mlx_array_item_bool_ptr(res, arr);
-}
-
-int mlx_array_item_uint8(uint8_t* res, const mlx_array arr) {
-    return mlx_array_item_uint8_ptr(res, arr);
-}
-
-int mlx_array_item_uint16(uint16_t* res, const mlx_array arr) {
-    return mlx_array_item_uint16_ptr(res, arr);
-}
-
-int mlx_array_item_uint32(uint32_t* res, const mlx_array arr) {
-    return mlx_array_item_uint32_ptr(res, arr);
-}
-
-int mlx_array_item_uint64(uint64_t* res, const mlx_array arr) {
-    return mlx_array_item_uint64_ptr(res, arr);
-}
-
-int mlx_array_item_int8(int8_t* res, const mlx_array arr) {
-    return mlx_array_item_int8_ptr(res, arr);
-}
-
-int mlx_array_item_int16(int16_t* res, const mlx_array arr) {
-    return mlx_array_item_int16_ptr(res, arr);
-}
-
-int mlx_array_item_int32(int32_t* res, const mlx_array arr) {
-    return mlx_array_item_int32_ptr(res, arr);
-}
-
-int mlx_array_item_int64(int64_t* res, const mlx_array arr) {
-    return mlx_array_item_int64_ptr(res, arr);
-}
-
-int mlx_array_item_float32(float* res, const mlx_array arr) {
-    return mlx_array_item_float32_ptr(res, arr);
-}
-
-int mlx_array_item_float64(double* res, const mlx_array arr) {
-    return mlx_array_item_float64_ptr(res, arr);
-}
-
-int mlx_array_item_complex64(float _Complex* res, const mlx_array arr) {
-    return mlx_array_item_complex64_ptr(res, arr);
-}
-
-#if defined(__aarch64__) || defined(_M_ARM64)
-int mlx_array_item_float16(float16_t* res, const mlx_array arr) {
-    return mlx_array_item_float16_ptr(res, arr);
-}
-#endif
-
-#if defined(__aarch64__) || defined(_M_ARM64)
-int mlx_array_item_bfloat16(bfloat16_t* res, const mlx_array arr) {
-    return mlx_array_item_bfloat16_ptr(res, arr);
-}
-#endif
-
-const bool* mlx_array_data_bool(const mlx_array arr) {
-    return mlx_array_data_bool_ptr(arr);
-}
-
-const uint8_t* mlx_array_data_uint8(const mlx_array arr) {
-    return mlx_array_data_uint8_ptr(arr);
-}
-
-const uint16_t* mlx_array_data_uint16(const mlx_array arr) {
-    return mlx_array_data_uint16_ptr(arr);
-}
-
-const uint32_t* mlx_array_data_uint32(const mlx_array arr) {
-    return mlx_array_data_uint32_ptr(arr);
-}
-
-const uint64_t* mlx_array_data_uint64(const mlx_array arr) {
-    return mlx_array_data_uint64_ptr(arr);
-}
-
-const int8_t* mlx_array_data_int8(const mlx_array arr) {
-    return mlx_array_data_int8_ptr(arr);
-}
-
-const int16_t* mlx_array_data_int16(const mlx_array arr) {
-    return mlx_array_data_int16_ptr(arr);
-}
-
-const int32_t* mlx_array_data_int32(const mlx_array arr) {
-    return mlx_array_data_int32_ptr(arr);
-}
-
-const int64_t* mlx_array_data_int64(const mlx_array arr) {
-    return mlx_array_data_int64_ptr(arr);
-}
-
-const float* mlx_array_data_float32(const mlx_array arr) {
-    return mlx_array_data_float32_ptr(arr);
-}
-
-const double* mlx_array_data_float64(const mlx_array arr) {
-    return mlx_array_data_float64_ptr(arr);
-}
-
-const float _Complex* mlx_array_data_complex64(const mlx_array arr) {
-    return mlx_array_data_complex64_ptr(arr);
-}
-
-#if defined(__aarch64__) || defined(_M_ARM64)
-const float16_t* mlx_array_data_float16(const mlx_array arr) {
-    return mlx_array_data_float16_ptr(arr);
-}
-#endif
-
-#if defined(__aarch64__) || defined(_M_ARM64)
-const bfloat16_t* mlx_array_data_bfloat16(const mlx_array arr) {
-    return mlx_array_data_bfloat16_ptr(arr);
-}
-#endif
-
-int _mlx_array_is_available(bool* res, const mlx_array arr) {
-    return _mlx_array_is_available_ptr(res, arr);
-}
-
-int _mlx_array_wait(const mlx_array arr) {
-    return _mlx_array_wait_ptr(arr);
-}
-
-int _mlx_array_is_contiguous(bool* res, const mlx_array arr) {
-    return _mlx_array_is_contiguous_ptr(res, arr);
-}
-
-int _mlx_array_is_row_contiguous(bool* res, const mlx_array arr) {
-    return _mlx_array_is_row_contiguous_ptr(res, arr);
-}
-
-int _mlx_array_is_col_contiguous(bool* res, const mlx_array arr) {
-    return _mlx_array_is_col_contiguous_ptr(res, arr);
-}
-
-mlx_closure mlx_closure_new(void) {
-    return mlx_closure_new_ptr();
-}
-
-int mlx_closure_free(mlx_closure cls) {
-    return mlx_closure_free_ptr(cls);
-}
-
-mlx_closure mlx_closure_new_func(int (*fun)(mlx_vector_array*, const mlx_vector_array)) {
-    return mlx_closure_new_func_ptr(fun);
-}
-
-mlx_closure mlx_closure_new_func_payload(int (*fun)(mlx_vector_array*, const mlx_vector_array, void*), void* payload, void (*dtor)(void*)) {
-    return mlx_closure_new_func_payload_ptr(fun, payload, dtor);
-}
-
-int mlx_closure_set(mlx_closure* cls, const mlx_closure src) {
-    return mlx_closure_set_ptr(cls, src);
-}
-
-int mlx_closure_apply(mlx_vector_array* res, mlx_closure cls, const mlx_vector_array input) {
-    return mlx_closure_apply_ptr(res, cls, input);
-}
-
-mlx_closure mlx_closure_new_unary(int (*fun)(mlx_array*, const mlx_array)) {
-    return mlx_closure_new_unary_ptr(fun);
-}
-
-mlx_closure_kwargs mlx_closure_kwargs_new(void) {
-    return mlx_closure_kwargs_new_ptr();
-}
-
-int mlx_closure_kwargs_free(mlx_closure_kwargs cls) {
-    return mlx_closure_kwargs_free_ptr(cls);
-}
-
-mlx_closure_kwargs mlx_closure_kwargs_new_func(int (*fun)( mlx_vector_array*, const mlx_vector_array, const mlx_map_string_to_array)) {
-    return mlx_closure_kwargs_new_func_ptr(fun);
-}
-
-mlx_closure_kwargs mlx_closure_kwargs_new_func_payload(int (*fun)( mlx_vector_array*, const mlx_vector_array, const mlx_map_string_to_array, void*), void* payload, void (*dtor)(void*)) {
-    return mlx_closure_kwargs_new_func_payload_ptr(fun, payload, dtor);
-}
-
-int mlx_closure_kwargs_set(mlx_closure_kwargs* cls, const mlx_closure_kwargs src) {
-    return mlx_closure_kwargs_set_ptr(cls, src);
-}
-
-int mlx_closure_kwargs_apply(mlx_vector_array* res, mlx_closure_kwargs cls, const mlx_vector_array input_0, const mlx_map_string_to_array input_1) {
-    return mlx_closure_kwargs_apply_ptr(res, cls, input_0, input_1);
-}
-
-mlx_closure_value_and_grad mlx_closure_value_and_grad_new(void) {
-    return mlx_closure_value_and_grad_new_ptr();
-}
-
-int mlx_closure_value_and_grad_free(mlx_closure_value_and_grad cls) {
-    return mlx_closure_value_and_grad_free_ptr(cls);
-}
-
-mlx_closure_value_and_grad mlx_closure_value_and_grad_new_func(int (*fun)(mlx_vector_array*, mlx_vector_array*, const mlx_vector_array)) {
-    return mlx_closure_value_and_grad_new_func_ptr(fun);
-}
-
-mlx_closure_value_and_grad mlx_closure_value_and_grad_new_func_payload(int (*fun)( mlx_vector_array*, mlx_vector_array*, const mlx_vector_array, void*), void* payload, void (*dtor)(void*)) {
-    return mlx_closure_value_and_grad_new_func_payload_ptr(fun, payload, dtor);
-}
-
-int mlx_closure_value_and_grad_set(mlx_closure_value_and_grad* cls, const mlx_closure_value_and_grad src) {
-    return mlx_closure_value_and_grad_set_ptr(cls, src);
-}
-
-int mlx_closure_value_and_grad_apply(mlx_vector_array* res_0, mlx_vector_array* res_1, mlx_closure_value_and_grad cls, const mlx_vector_array input) {
-    return mlx_closure_value_and_grad_apply_ptr(res_0, res_1, cls, input);
-}
-
-mlx_closure_custom mlx_closure_custom_new(void) {
-    return mlx_closure_custom_new_ptr();
-}
-
-int mlx_closure_custom_free(mlx_closure_custom cls) {
-    return mlx_closure_custom_free_ptr(cls);
-}
-
-mlx_closure_custom mlx_closure_custom_new_func(int (*fun)( mlx_vector_array*, const mlx_vector_array, const mlx_vector_array, const mlx_vector_array)) {
-    return mlx_closure_custom_new_func_ptr(fun);
-}
-
-mlx_closure_custom mlx_closure_custom_new_func_payload(int (*fun)( mlx_vector_array*, const mlx_vector_array, const mlx_vector_array, const mlx_vector_array, void*), void* payload, void (*dtor)(void*)) {
-    return mlx_closure_custom_new_func_payload_ptr(fun, payload, dtor);
-}
-
-int mlx_closure_custom_set(mlx_closure_custom* cls, const mlx_closure_custom src) {
-    return mlx_closure_custom_set_ptr(cls, src);
-}
-
-int mlx_closure_custom_apply(mlx_vector_array* res, mlx_closure_custom cls, const mlx_vector_array input_0, const mlx_vector_array input_1, const mlx_vector_array input_2) {
-    return mlx_closure_custom_apply_ptr(res, cls, input_0, input_1, input_2);
-}
-
-mlx_closure_custom_jvp mlx_closure_custom_jvp_new(void) {
-    return mlx_closure_custom_jvp_new_ptr();
-}
-
-int mlx_closure_custom_jvp_free(mlx_closure_custom_jvp cls) {
-    return mlx_closure_custom_jvp_free_ptr(cls);
-}
-
-mlx_closure_custom_jvp mlx_closure_custom_jvp_new_func(int (*fun)( mlx_vector_array*, const mlx_vector_array, const mlx_vector_array, const int*, size_t _num)) {
-    return mlx_closure_custom_jvp_new_func_ptr(fun);
-}
-
-mlx_closure_custom_jvp mlx_closure_custom_jvp_new_func_payload(int (*fun)( mlx_vector_array*, const mlx_vector_array, const mlx_vector_array, const int*, size_t _num, void*), void* payload, void (*dtor)(void*)) {
-    return mlx_closure_custom_jvp_new_func_payload_ptr(fun, payload, dtor);
-}
-
-int mlx_closure_custom_jvp_set(mlx_closure_custom_jvp* cls, const mlx_closure_custom_jvp src) {
-    return mlx_closure_custom_jvp_set_ptr(cls, src);
-}
-
-int mlx_closure_custom_jvp_apply(mlx_vector_array* res, mlx_closure_custom_jvp cls, const mlx_vector_array input_0, const mlx_vector_array input_1, const int* input_2, size_t input_2_num) {
-    return mlx_closure_custom_jvp_apply_ptr(res, cls, input_0, input_1, input_2, input_2_num);
-}
-
-mlx_closure_custom_vmap mlx_closure_custom_vmap_new(void) {
-    return mlx_closure_custom_vmap_new_ptr();
-}
-
-int mlx_closure_custom_vmap_free(mlx_closure_custom_vmap cls) {
-    return mlx_closure_custom_vmap_free_ptr(cls);
-}
-
-mlx_closure_custom_vmap mlx_closure_custom_vmap_new_func(int (*fun)( mlx_vector_array*, mlx_vector_int*, const mlx_vector_array, const int*, size_t _num)) {
-    return mlx_closure_custom_vmap_new_func_ptr(fun);
-}
-
-mlx_closure_custom_vmap mlx_closure_custom_vmap_new_func_payload(int (*fun)( mlx_vector_array*, mlx_vector_int*, const mlx_vector_array, const int*, size_t _num, void*), void* payload, void (*dtor)(void*)) {
-    return mlx_closure_custom_vmap_new_func_payload_ptr(fun, payload, dtor);
-}
-
-int mlx_closure_custom_vmap_set(mlx_closure_custom_vmap* cls, const mlx_closure_custom_vmap src) {
-    return mlx_closure_custom_vmap_set_ptr(cls, src);
-}
-
-int mlx_closure_custom_vmap_apply(mlx_vector_array* res_0, mlx_vector_int* res_1, mlx_closure_custom_vmap cls, const mlx_vector_array input_0, const int* input_1, size_t input_1_num) {
-    return mlx_closure_custom_vmap_apply_ptr(res_0, res_1, cls, input_0, input_1, input_1_num);
-}
-
-int mlx_compile(mlx_closure* res, const mlx_closure fun, bool shapeless) {
-    return mlx_compile_ptr(res, fun, shapeless);
-}
-
-int mlx_detail_compile(mlx_closure* res, const mlx_closure fun, uintptr_t fun_id, bool shapeless, const uint64_t* constants, size_t constants_num) {
-    return mlx_detail_compile_ptr(res, fun, fun_id, shapeless, constants, constants_num);
-}
-
-int mlx_detail_compile_clear_cache(void) {
-    return mlx_detail_compile_clear_cache_ptr();
-}
-
-int mlx_detail_compile_erase(uintptr_t fun_id) {
-    return mlx_detail_compile_erase_ptr(fun_id);
-}
-
-int mlx_disable_compile(void) {
-    return mlx_disable_compile_ptr();
-}
-
-int mlx_enable_compile(void) {
-    return mlx_enable_compile_ptr();
-}
-
-int mlx_set_compile_mode(mlx_compile_mode mode) {
-    return mlx_set_compile_mode_ptr(mode);
-}
-
-mlx_device mlx_device_new(void) {
-    return mlx_device_new_ptr();
-}
-
-mlx_device mlx_device_new_type(mlx_device_type type, int index) {
-    return mlx_device_new_type_ptr(type, index);
-}
-
-int mlx_device_free(mlx_device dev) {
-    return mlx_device_free_ptr(dev);
-}
-
-int mlx_device_set(mlx_device* dev, const mlx_device src) {
-    return mlx_device_set_ptr(dev, src);
-}
-
-int mlx_device_tostring(mlx_string* str, mlx_device dev) {
-    return mlx_device_tostring_ptr(str, dev);
-}
-
-bool mlx_device_equal(mlx_device lhs, mlx_device rhs) {
-    return mlx_device_equal_ptr(lhs, rhs);
-}
-
-int mlx_device_get_index(int* index, mlx_device dev) {
-    return mlx_device_get_index_ptr(index, dev);
-}
-
-int mlx_device_get_type(mlx_device_type* type, mlx_device dev) {
-    return mlx_device_get_type_ptr(type, dev);
-}
-
-int mlx_get_default_device(mlx_device* dev) {
-    return mlx_get_default_device_ptr(dev);
-}
-
-int mlx_set_default_device(mlx_device dev) {
-    return mlx_set_default_device_ptr(dev);
-}
-
-int mlx_distributed_all_gather(mlx_array* res, const mlx_array x, const mlx_distributed_group group , const mlx_stream S) {
-    return mlx_distributed_all_gather_ptr(res, x, group, S);
-}
-
-int mlx_distributed_all_max(mlx_array* res, const mlx_array x, const mlx_distributed_group group , const mlx_stream s) {
-    return mlx_distributed_all_max_ptr(res, x, group, s);
-}
-
-int mlx_distributed_all_min(mlx_array* res, const mlx_array x, const mlx_distributed_group group , const mlx_stream s) {
-    return mlx_distributed_all_min_ptr(res, x, group, s);
-}
-
-int mlx_distributed_all_sum(mlx_array* res, const mlx_array x, const mlx_distributed_group group , const mlx_stream s) {
-    return mlx_distributed_all_sum_ptr(res, x, group, s);
-}
-
-int mlx_distributed_recv(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, int src, const mlx_distributed_group group , const mlx_stream s) {
-    return mlx_distributed_recv_ptr(res, shape, shape_num, dtype, src, group, s);
-}
-
-int mlx_distributed_recv_like(mlx_array* res, const mlx_array x, int src, const mlx_distributed_group group , const mlx_stream s) {
-    return mlx_distributed_recv_like_ptr(res, x, src, group, s);
-}
-
-int mlx_distributed_send(mlx_array* res, const mlx_array x, int dst, const mlx_distributed_group group , const mlx_stream s) {
-    return mlx_distributed_send_ptr(res, x, dst, group, s);
-}
-
-int mlx_distributed_sum_scatter(mlx_array* res, const mlx_array x, const mlx_distributed_group group , const mlx_stream s) {
-    return mlx_distributed_sum_scatter_ptr(res, x, group, s);
-}
-
-int mlx_distributed_group_rank(mlx_distributed_group group) {
-    return mlx_distributed_group_rank_ptr(group);
-}
-
-int mlx_distributed_group_size(mlx_distributed_group group) {
-    return mlx_distributed_group_size_ptr(group);
-}
-
-mlx_distributed_group mlx_distributed_group_split(mlx_distributed_group group, int color, int key) {
-    return mlx_distributed_group_split_ptr(group, color, key);
-}
-
-bool mlx_distributed_is_available(void) {
-    return mlx_distributed_is_available_ptr();
-}
-
-mlx_distributed_group mlx_distributed_init(bool strict) {
-    return mlx_distributed_init_ptr(strict);
-}
-
-void mlx_set_error_handler(mlx_error_handler_func handler, void* data, void (*dtor)(void*)) {
-    mlx_set_error_handler_ptr(handler, data, dtor);
-}
-
-void _mlx_error(const char* file, const int line, const char* fmt, ...) {
-    _mlx_error_ptr(file, line, fmt);
-}
-
-int mlx_export_function(const char* file, const mlx_closure fun, const mlx_vector_array args, bool shapeless) {
-    return mlx_export_function_ptr(file, fun, args, shapeless);
-}
-
-int mlx_export_function_kwargs(const char* file, const mlx_closure_kwargs fun, const mlx_vector_array args, const mlx_map_string_to_array kwargs, bool shapeless) {
-    return mlx_export_function_kwargs_ptr(file, fun, args, kwargs, shapeless);
-}
-
-mlx_function_exporter mlx_function_exporter_new(const char* file, const mlx_closure fun, bool shapeless) {
-    return mlx_function_exporter_new_ptr(file, fun, shapeless);
-}
-
-int mlx_function_exporter_free(mlx_function_exporter xfunc) {
-    return mlx_function_exporter_free_ptr(xfunc);
-}
-
-int mlx_function_exporter_apply(const mlx_function_exporter xfunc, const mlx_vector_array args) {
-    return mlx_function_exporter_apply_ptr(xfunc, args);
-}
-
-int mlx_function_exporter_apply_kwargs(const mlx_function_exporter xfunc, const mlx_vector_array args, const mlx_map_string_to_array kwargs) {
-    return mlx_function_exporter_apply_kwargs_ptr(xfunc, args, kwargs);
-}
-
-mlx_imported_function mlx_imported_function_new(const char* file) {
-    return mlx_imported_function_new_ptr(file);
-}
-
-int mlx_imported_function_free(mlx_imported_function xfunc) {
-    return mlx_imported_function_free_ptr(xfunc);
-}
-
-int mlx_imported_function_apply(mlx_vector_array* res, const mlx_imported_function xfunc, const mlx_vector_array args) {
-    return mlx_imported_function_apply_ptr(res, xfunc, args);
-}
-
-int mlx_imported_function_apply_kwargs(mlx_vector_array* res, const mlx_imported_function xfunc, const mlx_vector_array args, const mlx_map_string_to_array kwargs) {
-    return mlx_imported_function_apply_kwargs_ptr(res, xfunc, args, kwargs);
-}
-
-mlx_fast_cuda_kernel_config mlx_fast_cuda_kernel_config_new(void) {
-    return mlx_fast_cuda_kernel_config_new_ptr();
-}
-
-void mlx_fast_cuda_kernel_config_free(mlx_fast_cuda_kernel_config cls) {
-    mlx_fast_cuda_kernel_config_free_ptr(cls);
-}
-
-int mlx_fast_cuda_kernel_config_add_output_arg(mlx_fast_cuda_kernel_config cls, const int* shape, size_t size, mlx_dtype dtype) {
-    return mlx_fast_cuda_kernel_config_add_output_arg_ptr(cls, shape, size, dtype);
-}
-
-int mlx_fast_cuda_kernel_config_set_grid(mlx_fast_cuda_kernel_config cls, int grid1, int grid2, int grid3) {
-    return mlx_fast_cuda_kernel_config_set_grid_ptr(cls, grid1, grid2, grid3);
-}
-
-int mlx_fast_cuda_kernel_config_set_thread_group(mlx_fast_cuda_kernel_config cls, int thread1, int thread2, int thread3) {
-    return mlx_fast_cuda_kernel_config_set_thread_group_ptr(cls, thread1, thread2, thread3);
-}
-
-int mlx_fast_cuda_kernel_config_set_init_value(mlx_fast_cuda_kernel_config cls, float value) {
-    return mlx_fast_cuda_kernel_config_set_init_value_ptr(cls, value);
-}
-
-int mlx_fast_cuda_kernel_config_set_verbose(mlx_fast_cuda_kernel_config cls, bool verbose) {
-    return mlx_fast_cuda_kernel_config_set_verbose_ptr(cls, verbose);
-}
-
-int mlx_fast_cuda_kernel_config_add_template_arg_dtype(mlx_fast_cuda_kernel_config cls, const char* name, mlx_dtype dtype) {
-    return mlx_fast_cuda_kernel_config_add_template_arg_dtype_ptr(cls, name, dtype);
-}
-
-int mlx_fast_cuda_kernel_config_add_template_arg_int(mlx_fast_cuda_kernel_config cls, const char* name, int value) {
-    return mlx_fast_cuda_kernel_config_add_template_arg_int_ptr(cls, name, value);
-}
-
-int mlx_fast_cuda_kernel_config_add_template_arg_bool(mlx_fast_cuda_kernel_config cls, const char* name, bool value) {
-    return mlx_fast_cuda_kernel_config_add_template_arg_bool_ptr(cls, name, value);
-}
-
-mlx_fast_cuda_kernel mlx_fast_cuda_kernel_new(const char* name, const mlx_vector_string input_names, const mlx_vector_string output_names, const char* source, const char* header, bool ensure_row_contiguous, int shared_memory) {
-    return mlx_fast_cuda_kernel_new_ptr(name, input_names, output_names, source, header, ensure_row_contiguous, shared_memory);
-}
-
-void mlx_fast_cuda_kernel_free(mlx_fast_cuda_kernel cls) {
-    mlx_fast_cuda_kernel_free_ptr(cls);
-}
-
-int mlx_fast_cuda_kernel_apply(mlx_vector_array* outputs, mlx_fast_cuda_kernel cls, const mlx_vector_array inputs, const mlx_fast_cuda_kernel_config config, const mlx_stream stream) {
-    return mlx_fast_cuda_kernel_apply_ptr(outputs, cls, inputs, config, stream);
-}
-
-int mlx_fast_layer_norm(mlx_array* res, const mlx_array x, const mlx_array weight , const mlx_array bias , float eps, const mlx_stream s) {
-    return mlx_fast_layer_norm_ptr(res, x, weight, bias, eps, s);
-}
-
-mlx_fast_metal_kernel_config mlx_fast_metal_kernel_config_new(void) {
-    return mlx_fast_metal_kernel_config_new_ptr();
-}
-
-void mlx_fast_metal_kernel_config_free(mlx_fast_metal_kernel_config cls) {
-    mlx_fast_metal_kernel_config_free_ptr(cls);
-}
-
-int mlx_fast_metal_kernel_config_add_output_arg(mlx_fast_metal_kernel_config cls, const int* shape, size_t size, mlx_dtype dtype) {
-    return mlx_fast_metal_kernel_config_add_output_arg_ptr(cls, shape, size, dtype);
-}
-
-int mlx_fast_metal_kernel_config_set_grid(mlx_fast_metal_kernel_config cls, int grid1, int grid2, int grid3) {
-    return mlx_fast_metal_kernel_config_set_grid_ptr(cls, grid1, grid2, grid3);
-}
-
-int mlx_fast_metal_kernel_config_set_thread_group(mlx_fast_metal_kernel_config cls, int thread1, int thread2, int thread3) {
-    return mlx_fast_metal_kernel_config_set_thread_group_ptr(cls, thread1, thread2, thread3);
-}
-
-int mlx_fast_metal_kernel_config_set_init_value(mlx_fast_metal_kernel_config cls, float value) {
-    return mlx_fast_metal_kernel_config_set_init_value_ptr(cls, value);
-}
-
-int mlx_fast_metal_kernel_config_set_verbose(mlx_fast_metal_kernel_config cls, bool verbose) {
-    return mlx_fast_metal_kernel_config_set_verbose_ptr(cls, verbose);
-}
-
-int mlx_fast_metal_kernel_config_add_template_arg_dtype(mlx_fast_metal_kernel_config cls, const char* name, mlx_dtype dtype) {
-    return mlx_fast_metal_kernel_config_add_template_arg_dtype_ptr(cls, name, dtype);
-}
-
-int mlx_fast_metal_kernel_config_add_template_arg_int(mlx_fast_metal_kernel_config cls, const char* name, int value) {
-    return mlx_fast_metal_kernel_config_add_template_arg_int_ptr(cls, name, value);
-}
-
-int mlx_fast_metal_kernel_config_add_template_arg_bool(mlx_fast_metal_kernel_config cls, const char* name, bool value) {
-    return mlx_fast_metal_kernel_config_add_template_arg_bool_ptr(cls, name, value);
-}
-
-mlx_fast_metal_kernel mlx_fast_metal_kernel_new(const char* name, const mlx_vector_string input_names, const mlx_vector_string output_names, const char* source, const char* header, bool ensure_row_contiguous, bool atomic_outputs) {
-    return mlx_fast_metal_kernel_new_ptr(name, input_names, output_names, source, header, ensure_row_contiguous, atomic_outputs);
-}
-
-void mlx_fast_metal_kernel_free(mlx_fast_metal_kernel cls) {
-    mlx_fast_metal_kernel_free_ptr(cls);
-}
-
-int mlx_fast_metal_kernel_apply(mlx_vector_array* outputs, mlx_fast_metal_kernel cls, const mlx_vector_array inputs, const mlx_fast_metal_kernel_config config, const mlx_stream stream) {
-    return mlx_fast_metal_kernel_apply_ptr(outputs, cls, inputs, config, stream);
-}
-
-int mlx_fast_rms_norm(mlx_array* res, const mlx_array x, const mlx_array weight , float eps, const mlx_stream s) {
-    return mlx_fast_rms_norm_ptr(res, x, weight, eps, s);
-}
-
-int mlx_fast_rope(mlx_array* res, const mlx_array x, int dims, bool traditional, mlx_optional_float base, float scale, int offset, const mlx_array freqs , const mlx_stream s) {
-    return mlx_fast_rope_ptr(res, x, dims, traditional, base, scale, offset, freqs, s);
-}
-
-int mlx_fast_rope_dynamic(mlx_array* res, const mlx_array x, int dims, bool traditional, mlx_optional_float base, float scale, const mlx_array offset, const mlx_array freqs , const mlx_stream s) {
-    return mlx_fast_rope_dynamic_ptr(res, x, dims, traditional, base, scale, offset, freqs, s);
-}
-
-int mlx_fast_scaled_dot_product_attention(mlx_array* res, const mlx_array queries, const mlx_array keys, const mlx_array values, float scale, const char* mask_mode, const mlx_array mask_arr , const mlx_array sinks , const mlx_stream s) {
-    return mlx_fast_scaled_dot_product_attention_ptr(res, queries, keys, values, scale, mask_mode, mask_arr, sinks, s);
-}
-
-int mlx_fft_fft(mlx_array* res, const mlx_array a, int n, int axis, const mlx_stream s) {
-    return mlx_fft_fft_ptr(res, a, n, axis, s);
-}
-
-int mlx_fft_fft2(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s) {
-    return mlx_fft_fft2_ptr(res, a, n, n_num, axes, axes_num, s);
-}
-
-int mlx_fft_fftn(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s) {
-    return mlx_fft_fftn_ptr(res, a, n, n_num, axes, axes_num, s);
-}
-
-int mlx_fft_fftshift(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, const mlx_stream s) {
-    return mlx_fft_fftshift_ptr(res, a, axes, axes_num, s);
-}
-
-int mlx_fft_ifft(mlx_array* res, const mlx_array a, int n, int axis, const mlx_stream s) {
-    return mlx_fft_ifft_ptr(res, a, n, axis, s);
-}
-
-int mlx_fft_ifft2(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s) {
-    return mlx_fft_ifft2_ptr(res, a, n, n_num, axes, axes_num, s);
-}
-
-int mlx_fft_ifftn(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s) {
-    return mlx_fft_ifftn_ptr(res, a, n, n_num, axes, axes_num, s);
-}
-
-int mlx_fft_ifftshift(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, const mlx_stream s) {
-    return mlx_fft_ifftshift_ptr(res, a, axes, axes_num, s);
-}
-
-int mlx_fft_irfft(mlx_array* res, const mlx_array a, int n, int axis, const mlx_stream s) {
-    return mlx_fft_irfft_ptr(res, a, n, axis, s);
-}
-
-int mlx_fft_irfft2(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s) {
-    return mlx_fft_irfft2_ptr(res, a, n, n_num, axes, axes_num, s);
-}
-
-int mlx_fft_irfftn(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s) {
-    return mlx_fft_irfftn_ptr(res, a, n, n_num, axes, axes_num, s);
-}
-
-int mlx_fft_rfft(mlx_array* res, const mlx_array a, int n, int axis, const mlx_stream s) {
-    return mlx_fft_rfft_ptr(res, a, n, axis, s);
-}
-
-int mlx_fft_rfft2(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s) {
-    return mlx_fft_rfft2_ptr(res, a, n, n_num, axes, axes_num, s);
-}
-
-int mlx_fft_rfftn(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s) {
-    return mlx_fft_rfftn_ptr(res, a, n, n_num, axes, axes_num, s);
-}
-
-int mlx_load_reader(mlx_array* res, mlx_io_reader in_stream, const mlx_stream s) {
-    return mlx_load_reader_ptr(res, in_stream, s);
-}
-
-int mlx_load(mlx_array* res, const char* file, const mlx_stream s) {
-    return mlx_load_ptr(res, file, s);
-}
-
-int mlx_load_safetensors_reader(mlx_map_string_to_array* res_0, mlx_map_string_to_string* res_1, mlx_io_reader in_stream, const mlx_stream s) {
-    return mlx_load_safetensors_reader_ptr(res_0, res_1, in_stream, s);
-}
-
-int mlx_load_safetensors(mlx_map_string_to_array* res_0, mlx_map_string_to_string* res_1, const char* file, const mlx_stream s) {
-    return mlx_load_safetensors_ptr(res_0, res_1, file, s);
-}
-
-int mlx_save_writer(mlx_io_writer out_stream, const mlx_array a) {
-    return mlx_save_writer_ptr(out_stream, a);
-}
-
-int mlx_save(const char* file, const mlx_array a) {
-    return mlx_save_ptr(file, a);
-}
-
-int mlx_save_safetensors_writer(mlx_io_writer in_stream, const mlx_map_string_to_array param, const mlx_map_string_to_string metadata) {
-    return mlx_save_safetensors_writer_ptr(in_stream, param, metadata);
-}
-
-int mlx_save_safetensors(const char* file, const mlx_map_string_to_array param, const mlx_map_string_to_string metadata) {
-    return mlx_save_safetensors_ptr(file, param, metadata);
-}
-
-mlx_io_reader mlx_io_reader_new(void* desc, mlx_io_vtable vtable) {
-    return mlx_io_reader_new_ptr(desc, vtable);
-}
-
-int mlx_io_reader_descriptor(void** desc_, mlx_io_reader io) {
-    return mlx_io_reader_descriptor_ptr(desc_, io);
-}
-
-int mlx_io_reader_tostring(mlx_string* str_, mlx_io_reader io) {
-    return mlx_io_reader_tostring_ptr(str_, io);
-}
-
-int mlx_io_reader_free(mlx_io_reader io) {
-    return mlx_io_reader_free_ptr(io);
-}
-
-mlx_io_writer mlx_io_writer_new(void* desc, mlx_io_vtable vtable) {
-    return mlx_io_writer_new_ptr(desc, vtable);
-}
-
-int mlx_io_writer_descriptor(void** desc_, mlx_io_writer io) {
-    return mlx_io_writer_descriptor_ptr(desc_, io);
-}
-
-int mlx_io_writer_tostring(mlx_string* str_, mlx_io_writer io) {
-    return mlx_io_writer_tostring_ptr(str_, io);
-}
-
-int mlx_io_writer_free(mlx_io_writer io) {
-    return mlx_io_writer_free_ptr(io);
-}
-
-int mlx_linalg_cholesky(mlx_array* res, const mlx_array a, bool upper, const mlx_stream s) {
-    return mlx_linalg_cholesky_ptr(res, a, upper, s);
-}
-
-int mlx_linalg_cholesky_inv(mlx_array* res, const mlx_array a, bool upper, const mlx_stream s) {
-    return mlx_linalg_cholesky_inv_ptr(res, a, upper, s);
-}
-
-int mlx_linalg_cross(mlx_array* res, const mlx_array a, const mlx_array b, int axis, const mlx_stream s) {
-    return mlx_linalg_cross_ptr(res, a, b, axis, s);
-}
-
-int mlx_linalg_eig(mlx_array* res_0, mlx_array* res_1, const mlx_array a, const mlx_stream s) {
-    return mlx_linalg_eig_ptr(res_0, res_1, a, s);
-}
-
-int mlx_linalg_eigh(mlx_array* res_0, mlx_array* res_1, const mlx_array a, const char* UPLO, const mlx_stream s) {
-    return mlx_linalg_eigh_ptr(res_0, res_1, a, UPLO, s);
-}
-
-int mlx_linalg_eigvals(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_linalg_eigvals_ptr(res, a, s);
-}
-
-int mlx_linalg_eigvalsh(mlx_array* res, const mlx_array a, const char* UPLO, const mlx_stream s) {
-    return mlx_linalg_eigvalsh_ptr(res, a, UPLO, s);
-}
-
-int mlx_linalg_inv(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_linalg_inv_ptr(res, a, s);
-}
-
-int mlx_linalg_lu(mlx_vector_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_linalg_lu_ptr(res, a, s);
-}
-
-int mlx_linalg_lu_factor(mlx_array* res_0, mlx_array* res_1, const mlx_array a, const mlx_stream s) {
-    return mlx_linalg_lu_factor_ptr(res_0, res_1, a, s);
-}
-
-int mlx_linalg_norm(mlx_array* res, const mlx_array a, double ord, const int* axis , size_t axis_num, bool keepdims, const mlx_stream s) {
-    return mlx_linalg_norm_ptr(res, a, ord, axis, axis_num, keepdims, s);
-}
-
-int mlx_linalg_norm_matrix(mlx_array* res, const mlx_array a, const char* ord, const int* axis , size_t axis_num, bool keepdims, const mlx_stream s) {
-    return mlx_linalg_norm_matrix_ptr(res, a, ord, axis, axis_num, keepdims, s);
-}
-
-int mlx_linalg_norm_l2(mlx_array* res, const mlx_array a, const int* axis , size_t axis_num, bool keepdims, const mlx_stream s) {
-    return mlx_linalg_norm_l2_ptr(res, a, axis, axis_num, keepdims, s);
-}
-
-int mlx_linalg_pinv(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_linalg_pinv_ptr(res, a, s);
-}
-
-int mlx_linalg_qr(mlx_array* res_0, mlx_array* res_1, const mlx_array a, const mlx_stream s) {
-    return mlx_linalg_qr_ptr(res_0, res_1, a, s);
-}
-
-int mlx_linalg_solve(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_linalg_solve_ptr(res, a, b, s);
-}
-
-int mlx_linalg_solve_triangular(mlx_array* res, const mlx_array a, const mlx_array b, bool upper, const mlx_stream s) {
-    return mlx_linalg_solve_triangular_ptr(res, a, b, upper, s);
-}
-
-int mlx_linalg_svd(mlx_vector_array* res, const mlx_array a, bool compute_uv, const mlx_stream s) {
-    return mlx_linalg_svd_ptr(res, a, compute_uv, s);
-}
-
-int mlx_linalg_tri_inv(mlx_array* res, const mlx_array a, bool upper, const mlx_stream s) {
-    return mlx_linalg_tri_inv_ptr(res, a, upper, s);
-}
-
-mlx_map_string_to_array mlx_map_string_to_array_new(void) {
-    return mlx_map_string_to_array_new_ptr();
-}
-
-int mlx_map_string_to_array_set(mlx_map_string_to_array* map, const mlx_map_string_to_array src) {
-    return mlx_map_string_to_array_set_ptr(map, src);
-}
-
-int mlx_map_string_to_array_free(mlx_map_string_to_array map) {
-    return mlx_map_string_to_array_free_ptr(map);
-}
-
-int mlx_map_string_to_array_insert(mlx_map_string_to_array map, const char* key, const mlx_array value) {
-    return mlx_map_string_to_array_insert_ptr(map, key, value);
-}
-
-int mlx_map_string_to_array_get(mlx_array* value, const mlx_map_string_to_array map, const char* key) {
-    return mlx_map_string_to_array_get_ptr(value, map, key);
-}
-
-mlx_map_string_to_array_iterator mlx_map_string_to_array_iterator_new(mlx_map_string_to_array map) {
-    return mlx_map_string_to_array_iterator_new_ptr(map);
-}
-
-int mlx_map_string_to_array_iterator_free(mlx_map_string_to_array_iterator it) {
-    return mlx_map_string_to_array_iterator_free_ptr(it);
-}
-
-int mlx_map_string_to_array_iterator_next(const char** key, mlx_array* value, mlx_map_string_to_array_iterator it) {
-    return mlx_map_string_to_array_iterator_next_ptr(key, value, it);
-}
-
-mlx_map_string_to_string mlx_map_string_to_string_new(void) {
-    return mlx_map_string_to_string_new_ptr();
-}
-
-int mlx_map_string_to_string_set(mlx_map_string_to_string* map, const mlx_map_string_to_string src) {
-    return mlx_map_string_to_string_set_ptr(map, src);
-}
-
-int mlx_map_string_to_string_free(mlx_map_string_to_string map) {
-    return mlx_map_string_to_string_free_ptr(map);
-}
-
-int mlx_map_string_to_string_insert(mlx_map_string_to_string map, const char* key, const char* value) {
-    return mlx_map_string_to_string_insert_ptr(map, key, value);
-}
-
-int mlx_map_string_to_string_get(const char** value, const mlx_map_string_to_string map, const char* key) {
-    return mlx_map_string_to_string_get_ptr(value, map, key);
-}
-
-mlx_map_string_to_string_iterator mlx_map_string_to_string_iterator_new(mlx_map_string_to_string map) {
-    return mlx_map_string_to_string_iterator_new_ptr(map);
-}
-
-int mlx_map_string_to_string_iterator_free(mlx_map_string_to_string_iterator it) {
-    return mlx_map_string_to_string_iterator_free_ptr(it);
-}
-
-int mlx_map_string_to_string_iterator_next(const char** key, const char** value, mlx_map_string_to_string_iterator it) {
-    return mlx_map_string_to_string_iterator_next_ptr(key, value, it);
-}
-
-int mlx_clear_cache(void) {
-    return mlx_clear_cache_ptr();
-}
-
-int mlx_get_active_memory(size_t* res) {
-    return mlx_get_active_memory_ptr(res);
-}
-
-int mlx_get_cache_memory(size_t* res) {
-    return mlx_get_cache_memory_ptr(res);
-}
-
-int mlx_get_memory_limit(size_t* res) {
-    return mlx_get_memory_limit_ptr(res);
-}
-
-int mlx_get_peak_memory(size_t* res) {
-    return mlx_get_peak_memory_ptr(res);
-}
-
-int mlx_reset_peak_memory(void) {
-    return mlx_reset_peak_memory_ptr();
-}
-
-int mlx_set_cache_limit(size_t* res, size_t limit) {
-    return mlx_set_cache_limit_ptr(res, limit);
-}
-
-int mlx_set_memory_limit(size_t* res, size_t limit) {
-    return mlx_set_memory_limit_ptr(res, limit);
-}
-
-int mlx_set_wired_limit(size_t* res, size_t limit) {
-    return mlx_set_wired_limit_ptr(res, limit);
-}
-
-mlx_metal_device_info_t mlx_metal_device_info(void) {
-    return mlx_metal_device_info_ptr();
-}
-
-int mlx_metal_is_available(bool* res) {
-    return mlx_metal_is_available_ptr(res);
-}
-
-int mlx_metal_start_capture(const char* path) {
-    return mlx_metal_start_capture_ptr(path);
-}
-
-int mlx_metal_stop_capture(void) {
-    return mlx_metal_stop_capture_ptr();
-}
-
-int mlx_abs(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_abs_ptr(res, a, s);
-}
-
-int mlx_add(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_add_ptr(res, a, b, s);
-}
-
-int mlx_addmm(mlx_array* res, const mlx_array c, const mlx_array a, const mlx_array b, float alpha, float beta, const mlx_stream s) {
-    return mlx_addmm_ptr(res, c, a, b, alpha, beta, s);
-}
-
-int mlx_all_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s) {
-    return mlx_all_axes_ptr(res, a, axes, axes_num, keepdims, s);
-}
-
-int mlx_all_axis(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s) {
-    return mlx_all_axis_ptr(res, a, axis, keepdims, s);
-}
-
-int mlx_all(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s) {
-    return mlx_all_ptr(res, a, keepdims, s);
-}
-
-int mlx_allclose(mlx_array* res, const mlx_array a, const mlx_array b, double rtol, double atol, bool equal_nan, const mlx_stream s) {
-    return mlx_allclose_ptr(res, a, b, rtol, atol, equal_nan, s);
-}
-
-int mlx_any_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s) {
-    return mlx_any_axes_ptr(res, a, axes, axes_num, keepdims, s);
-}
-
-int mlx_any_axis(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s) {
-    return mlx_any_axis_ptr(res, a, axis, keepdims, s);
-}
-
-int mlx_any(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s) {
-    return mlx_any_ptr(res, a, keepdims, s);
-}
-
-int mlx_arange(mlx_array* res, double start, double stop, double step, mlx_dtype dtype, const mlx_stream s) {
-    return mlx_arange_ptr(res, start, stop, step, dtype, s);
-}
-
-int mlx_arccos(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_arccos_ptr(res, a, s);
-}
-
-int mlx_arccosh(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_arccosh_ptr(res, a, s);
-}
-
-int mlx_arcsin(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_arcsin_ptr(res, a, s);
-}
-
-int mlx_arcsinh(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_arcsinh_ptr(res, a, s);
-}
-
-int mlx_arctan(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_arctan_ptr(res, a, s);
-}
-
-int mlx_arctan2(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_arctan2_ptr(res, a, b, s);
-}
-
-int mlx_arctanh(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_arctanh_ptr(res, a, s);
-}
-
-int mlx_argmax_axis(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s) {
-    return mlx_argmax_axis_ptr(res, a, axis, keepdims, s);
-}
-
-int mlx_argmax(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s) {
-    return mlx_argmax_ptr(res, a, keepdims, s);
-}
-
-int mlx_argmin_axis(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s) {
-    return mlx_argmin_axis_ptr(res, a, axis, keepdims, s);
-}
-
-int mlx_argmin(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s) {
-    return mlx_argmin_ptr(res, a, keepdims, s);
-}
-
-int mlx_argpartition_axis(mlx_array* res, const mlx_array a, int kth, int axis, const mlx_stream s) {
-    return mlx_argpartition_axis_ptr(res, a, kth, axis, s);
-}
-
-int mlx_argpartition(mlx_array* res, const mlx_array a, int kth, const mlx_stream s) {
-    return mlx_argpartition_ptr(res, a, kth, s);
-}
-
-int mlx_argsort_axis(mlx_array* res, const mlx_array a, int axis, const mlx_stream s) {
-    return mlx_argsort_axis_ptr(res, a, axis, s);
-}
-
-int mlx_argsort(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_argsort_ptr(res, a, s);
-}
-
-int mlx_array_equal(mlx_array* res, const mlx_array a, const mlx_array b, bool equal_nan, const mlx_stream s) {
-    return mlx_array_equal_ptr(res, a, b, equal_nan, s);
-}
-
-int mlx_as_strided(mlx_array* res, const mlx_array a, const int* shape, size_t shape_num, const int64_t* strides, size_t strides_num, size_t offset, const mlx_stream s) {
-    return mlx_as_strided_ptr(res, a, shape, shape_num, strides, strides_num, offset, s);
-}
-
-int mlx_astype(mlx_array* res, const mlx_array a, mlx_dtype dtype, const mlx_stream s) {
-    return mlx_astype_ptr(res, a, dtype, s);
-}
-
-int mlx_atleast_1d(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_atleast_1d_ptr(res, a, s);
-}
-
-int mlx_atleast_2d(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_atleast_2d_ptr(res, a, s);
-}
-
-int mlx_atleast_3d(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_atleast_3d_ptr(res, a, s);
-}
-
-int mlx_bitwise_and(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_bitwise_and_ptr(res, a, b, s);
-}
-
-int mlx_bitwise_invert(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_bitwise_invert_ptr(res, a, s);
-}
-
-int mlx_bitwise_or(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_bitwise_or_ptr(res, a, b, s);
-}
-
-int mlx_bitwise_xor(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_bitwise_xor_ptr(res, a, b, s);
-}
-
-int mlx_block_masked_mm(mlx_array* res, const mlx_array a, const mlx_array b, int block_size, const mlx_array mask_out , const mlx_array mask_lhs , const mlx_array mask_rhs , const mlx_stream s) {
-    return mlx_block_masked_mm_ptr(res, a, b, block_size, mask_out, mask_lhs, mask_rhs, s);
-}
-
-int mlx_broadcast_arrays(mlx_vector_array* res, const mlx_vector_array inputs, const mlx_stream s) {
-    return mlx_broadcast_arrays_ptr(res, inputs, s);
-}
-
-int mlx_broadcast_to(mlx_array* res, const mlx_array a, const int* shape, size_t shape_num, const mlx_stream s) {
-    return mlx_broadcast_to_ptr(res, a, shape, shape_num, s);
-}
-
-int mlx_ceil(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_ceil_ptr(res, a, s);
-}
-
-int mlx_clip(mlx_array* res, const mlx_array a, const mlx_array a_min , const mlx_array a_max , const mlx_stream s) {
-    return mlx_clip_ptr(res, a, a_min, a_max, s);
-}
-
-int mlx_concatenate_axis(mlx_array* res, const mlx_vector_array arrays, int axis, const mlx_stream s) {
-    return mlx_concatenate_axis_ptr(res, arrays, axis, s);
-}
-
-int mlx_concatenate(mlx_array* res, const mlx_vector_array arrays, const mlx_stream s) {
-    return mlx_concatenate_ptr(res, arrays, s);
-}
-
-int mlx_conjugate(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_conjugate_ptr(res, a, s);
-}
-
-int mlx_contiguous(mlx_array* res, const mlx_array a, bool allow_col_major, const mlx_stream s) {
-    return mlx_contiguous_ptr(res, a, allow_col_major, s);
-}
-
-int mlx_conv1d(mlx_array* res, const mlx_array input, const mlx_array weight, int stride, int padding, int dilation, int groups, const mlx_stream s) {
-    return mlx_conv1d_ptr(res, input, weight, stride, padding, dilation, groups, s);
-}
-
-int mlx_conv2d(mlx_array* res, const mlx_array input, const mlx_array weight, int stride_0, int stride_1, int padding_0, int padding_1, int dilation_0, int dilation_1, int groups, const mlx_stream s) {
-    return mlx_conv2d_ptr(res, input, weight, stride_0, stride_1, padding_0, padding_1, dilation_0, dilation_1, groups, s);
-}
-
-int mlx_conv3d(mlx_array* res, const mlx_array input, const mlx_array weight, int stride_0, int stride_1, int stride_2, int padding_0, int padding_1, int padding_2, int dilation_0, int dilation_1, int dilation_2, int groups, const mlx_stream s) {
-    return mlx_conv3d_ptr(res, input, weight, stride_0, stride_1, stride_2, padding_0, padding_1, padding_2, dilation_0, dilation_1, dilation_2, groups, s);
-}
-
-int mlx_conv_general(mlx_array* res, const mlx_array input, const mlx_array weight, const int* stride, size_t stride_num, const int* padding_lo, size_t padding_lo_num, const int* padding_hi, size_t padding_hi_num, const int* kernel_dilation, size_t kernel_dilation_num, const int* input_dilation, size_t input_dilation_num, int groups, bool flip, const mlx_stream s) {
-    return mlx_conv_general_ptr(res, input, weight, stride, stride_num, padding_lo, padding_lo_num, padding_hi, padding_hi_num, kernel_dilation, kernel_dilation_num, input_dilation, input_dilation_num, groups, flip, s);
-}
-
-int mlx_conv_transpose1d(mlx_array* res, const mlx_array input, const mlx_array weight, int stride, int padding, int dilation, int output_padding, int groups, const mlx_stream s) {
-    return mlx_conv_transpose1d_ptr(res, input, weight, stride, padding, dilation, output_padding, groups, s);
-}
-
-int mlx_conv_transpose2d(mlx_array* res, const mlx_array input, const mlx_array weight, int stride_0, int stride_1, int padding_0, int padding_1, int dilation_0, int dilation_1, int output_padding_0, int output_padding_1, int groups, const mlx_stream s) {
-    return mlx_conv_transpose2d_ptr(res, input, weight, stride_0, stride_1, padding_0, padding_1, dilation_0, dilation_1, output_padding_0, output_padding_1, groups, s);
-}
-
-int mlx_conv_transpose3d(mlx_array* res, const mlx_array input, const mlx_array weight, int stride_0, int stride_1, int stride_2, int padding_0, int padding_1, int padding_2, int dilation_0, int dilation_1, int dilation_2, int output_padding_0, int output_padding_1, int output_padding_2, int groups, const mlx_stream s) {
-    return mlx_conv_transpose3d_ptr(res, input, weight, stride_0, stride_1, stride_2, padding_0, padding_1, padding_2, dilation_0, dilation_1, dilation_2, output_padding_0, output_padding_1, output_padding_2, groups, s);
-}
-
-int mlx_copy(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_copy_ptr(res, a, s);
-}
-
-int mlx_cos(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_cos_ptr(res, a, s);
-}
-
-int mlx_cosh(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_cosh_ptr(res, a, s);
-}
-
-int mlx_cummax(mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s) {
-    return mlx_cummax_ptr(res, a, axis, reverse, inclusive, s);
-}
-
-int mlx_cummin(mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s) {
-    return mlx_cummin_ptr(res, a, axis, reverse, inclusive, s);
-}
-
-int mlx_cumprod(mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s) {
-    return mlx_cumprod_ptr(res, a, axis, reverse, inclusive, s);
-}
-
-int mlx_cumsum(mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s) {
-    return mlx_cumsum_ptr(res, a, axis, reverse, inclusive, s);
-}
-
-int mlx_degrees(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_degrees_ptr(res, a, s);
-}
-
-int mlx_depends(mlx_vector_array* res, const mlx_vector_array inputs, const mlx_vector_array dependencies) {
-    return mlx_depends_ptr(res, inputs, dependencies);
-}
-
-int mlx_dequantize(mlx_array* res, const mlx_array w, const mlx_array scales, const mlx_array biases , mlx_optional_int group_size, mlx_optional_int bits, const char* mode, mlx_optional_dtype dtype, const mlx_stream s) {
-    return mlx_dequantize_ptr(res, w, scales, biases, group_size, bits, mode, dtype, s);
-}
-
-int mlx_diag(mlx_array* res, const mlx_array a, int k, const mlx_stream s) {
-    return mlx_diag_ptr(res, a, k, s);
-}
-
-int mlx_diagonal(mlx_array* res, const mlx_array a, int offset, int axis1, int axis2, const mlx_stream s) {
-    return mlx_diagonal_ptr(res, a, offset, axis1, axis2, s);
-}
-
-int mlx_divide(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_divide_ptr(res, a, b, s);
-}
-
-int mlx_divmod(mlx_vector_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_divmod_ptr(res, a, b, s);
-}
-
-int mlx_einsum(mlx_array* res, const char* subscripts, const mlx_vector_array operands, const mlx_stream s) {
-    return mlx_einsum_ptr(res, subscripts, operands, s);
-}
-
-int mlx_equal(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_equal_ptr(res, a, b, s);
-}
-
-int mlx_erf(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_erf_ptr(res, a, s);
-}
-
-int mlx_erfinv(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_erfinv_ptr(res, a, s);
-}
-
-int mlx_exp(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_exp_ptr(res, a, s);
-}
-
-int mlx_expand_dims_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, const mlx_stream s) {
-    return mlx_expand_dims_axes_ptr(res, a, axes, axes_num, s);
-}
-
-int mlx_expand_dims(mlx_array* res, const mlx_array a, int axis, const mlx_stream s) {
-    return mlx_expand_dims_ptr(res, a, axis, s);
-}
-
-int mlx_expm1(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_expm1_ptr(res, a, s);
-}
-
-int mlx_eye(mlx_array* res, int n, int m, int k, mlx_dtype dtype, const mlx_stream s) {
-    return mlx_eye_ptr(res, n, m, k, dtype, s);
-}
-
-int mlx_flatten(mlx_array* res, const mlx_array a, int start_axis, int end_axis, const mlx_stream s) {
-    return mlx_flatten_ptr(res, a, start_axis, end_axis, s);
-}
-
-int mlx_floor(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_floor_ptr(res, a, s);
-}
-
-int mlx_floor_divide(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_floor_divide_ptr(res, a, b, s);
-}
-
-int mlx_from_fp8(mlx_array* res, const mlx_array x, mlx_dtype dtype, const mlx_stream s) {
-    return mlx_from_fp8_ptr(res, x, dtype, s);
-}
-
-int mlx_full(mlx_array* res, const int* shape, size_t shape_num, const mlx_array vals, mlx_dtype dtype, const mlx_stream s) {
-    return mlx_full_ptr(res, shape, shape_num, vals, dtype, s);
-}
-
-int mlx_full_like(mlx_array* res, const mlx_array a, const mlx_array vals, mlx_dtype dtype, const mlx_stream s) {
-    return mlx_full_like_ptr(res, a, vals, dtype, s);
-}
-
-int mlx_gather(mlx_array* res, const mlx_array a, const mlx_vector_array indices, const int* axes, size_t axes_num, const int* slice_sizes, size_t slice_sizes_num, const mlx_stream s) {
-    return mlx_gather_ptr(res, a, indices, axes, axes_num, slice_sizes, slice_sizes_num, s);
-}
-
-int mlx_gather_single(mlx_array* res, const mlx_array a, const mlx_array indices, int axis, const int* slice_sizes, size_t slice_sizes_num, const mlx_stream s) {
-    return mlx_gather_single_ptr(res, a, indices, axis, slice_sizes, slice_sizes_num, s);
-}
-
-int mlx_gather_mm(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_array lhs_indices , const mlx_array rhs_indices , bool sorted_indices, const mlx_stream s) {
-    return mlx_gather_mm_ptr(res, a, b, lhs_indices, rhs_indices, sorted_indices, s);
-}
-
-int mlx_gather_qmm(mlx_array* res, const mlx_array x, const mlx_array w, const mlx_array scales, const mlx_array biases , const mlx_array lhs_indices , const mlx_array rhs_indices , bool transpose, mlx_optional_int group_size, mlx_optional_int bits, const char* mode, bool sorted_indices, const mlx_stream s) {
-    return mlx_gather_qmm_ptr(res, x, w, scales, biases, lhs_indices, rhs_indices, transpose, group_size, bits, mode, sorted_indices, s);
-}
-
-int mlx_greater(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_greater_ptr(res, a, b, s);
-}
-
-int mlx_greater_equal(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_greater_equal_ptr(res, a, b, s);
-}
-
-int mlx_hadamard_transform(mlx_array* res, const mlx_array a, mlx_optional_float scale, const mlx_stream s) {
-    return mlx_hadamard_transform_ptr(res, a, scale, s);
-}
-
-int mlx_identity(mlx_array* res, int n, mlx_dtype dtype, const mlx_stream s) {
-    return mlx_identity_ptr(res, n, dtype, s);
-}
-
-int mlx_imag(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_imag_ptr(res, a, s);
-}
-
-int mlx_inner(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_inner_ptr(res, a, b, s);
-}
-
-int mlx_isclose(mlx_array* res, const mlx_array a, const mlx_array b, double rtol, double atol, bool equal_nan, const mlx_stream s) {
-    return mlx_isclose_ptr(res, a, b, rtol, atol, equal_nan, s);
-}
-
-int mlx_isfinite(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_isfinite_ptr(res, a, s);
-}
-
-int mlx_isinf(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_isinf_ptr(res, a, s);
-}
-
-int mlx_isnan(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_isnan_ptr(res, a, s);
-}
-
-int mlx_isneginf(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_isneginf_ptr(res, a, s);
-}
-
-int mlx_isposinf(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_isposinf_ptr(res, a, s);
-}
-
-int mlx_kron(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_kron_ptr(res, a, b, s);
-}
-
-int mlx_left_shift(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_left_shift_ptr(res, a, b, s);
-}
-
-int mlx_less(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_less_ptr(res, a, b, s);
-}
-
-int mlx_less_equal(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_less_equal_ptr(res, a, b, s);
-}
-
-int mlx_linspace(mlx_array* res, double start, double stop, int num, mlx_dtype dtype, const mlx_stream s) {
-    return mlx_linspace_ptr(res, start, stop, num, dtype, s);
-}
-
-int mlx_log(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_log_ptr(res, a, s);
-}
-
-int mlx_log10(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_log10_ptr(res, a, s);
-}
-
-int mlx_log1p(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_log1p_ptr(res, a, s);
-}
-
-int mlx_log2(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_log2_ptr(res, a, s);
-}
-
-int mlx_logaddexp(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_logaddexp_ptr(res, a, b, s);
-}
-
-int mlx_logcumsumexp(mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s) {
-    return mlx_logcumsumexp_ptr(res, a, axis, reverse, inclusive, s);
-}
-
-int mlx_logical_and(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_logical_and_ptr(res, a, b, s);
-}
-
-int mlx_logical_not(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_logical_not_ptr(res, a, s);
-}
-
-int mlx_logical_or(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_logical_or_ptr(res, a, b, s);
-}
-
-int mlx_logsumexp_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s) {
-    return mlx_logsumexp_axes_ptr(res, a, axes, axes_num, keepdims, s);
-}
-
-int mlx_logsumexp_axis(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s) {
-    return mlx_logsumexp_axis_ptr(res, a, axis, keepdims, s);
-}
-
-int mlx_logsumexp(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s) {
-    return mlx_logsumexp_ptr(res, a, keepdims, s);
-}
-
-int mlx_masked_scatter(mlx_array* res, const mlx_array a, const mlx_array mask, const mlx_array src, const mlx_stream s) {
-    return mlx_masked_scatter_ptr(res, a, mask, src, s);
-}
-
-int mlx_matmul(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_matmul_ptr(res, a, b, s);
-}
-
-int mlx_max_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s) {
-    return mlx_max_axes_ptr(res, a, axes, axes_num, keepdims, s);
-}
-
-int mlx_max_axis(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s) {
-    return mlx_max_axis_ptr(res, a, axis, keepdims, s);
-}
-
-int mlx_max(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s) {
-    return mlx_max_ptr(res, a, keepdims, s);
-}
-
-int mlx_maximum(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_maximum_ptr(res, a, b, s);
-}
-
-int mlx_mean_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s) {
-    return mlx_mean_axes_ptr(res, a, axes, axes_num, keepdims, s);
-}
-
-int mlx_mean_axis(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s) {
-    return mlx_mean_axis_ptr(res, a, axis, keepdims, s);
-}
-
-int mlx_mean(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s) {
-    return mlx_mean_ptr(res, a, keepdims, s);
-}
-
-int mlx_median(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s) {
-    return mlx_median_ptr(res, a, axes, axes_num, keepdims, s);
-}
-
-int mlx_meshgrid(mlx_vector_array* res, const mlx_vector_array arrays, bool sparse, const char* indexing, const mlx_stream s) {
-    return mlx_meshgrid_ptr(res, arrays, sparse, indexing, s);
-}
-
-int mlx_min_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s) {
-    return mlx_min_axes_ptr(res, a, axes, axes_num, keepdims, s);
-}
-
-int mlx_min_axis(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s) {
-    return mlx_min_axis_ptr(res, a, axis, keepdims, s);
-}
-
-int mlx_min(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s) {
-    return mlx_min_ptr(res, a, keepdims, s);
-}
-
-int mlx_minimum(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_minimum_ptr(res, a, b, s);
-}
-
-int mlx_moveaxis(mlx_array* res, const mlx_array a, int source, int destination, const mlx_stream s) {
-    return mlx_moveaxis_ptr(res, a, source, destination, s);
-}
-
-int mlx_multiply(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_multiply_ptr(res, a, b, s);
-}
-
-int mlx_nan_to_num(mlx_array* res, const mlx_array a, float nan, mlx_optional_float posinf, mlx_optional_float neginf, const mlx_stream s) {
-    return mlx_nan_to_num_ptr(res, a, nan, posinf, neginf, s);
-}
-
-int mlx_negative(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_negative_ptr(res, a, s);
-}
-
-int mlx_not_equal(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_not_equal_ptr(res, a, b, s);
-}
-
-int mlx_number_of_elements(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool inverted, mlx_dtype dtype, const mlx_stream s) {
-    return mlx_number_of_elements_ptr(res, a, axes, axes_num, inverted, dtype, s);
-}
-
-int mlx_ones(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_stream s) {
-    return mlx_ones_ptr(res, shape, shape_num, dtype, s);
-}
-
-int mlx_ones_like(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_ones_like_ptr(res, a, s);
-}
-
-int mlx_outer(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_outer_ptr(res, a, b, s);
-}
-
-int mlx_pad(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, const int* low_pad_size, size_t low_pad_size_num, const int* high_pad_size, size_t high_pad_size_num, const mlx_array pad_value, const char* mode, const mlx_stream s) {
-    return mlx_pad_ptr(res, a, axes, axes_num, low_pad_size, low_pad_size_num, high_pad_size, high_pad_size_num, pad_value, mode, s);
-}
-
-int mlx_pad_symmetric(mlx_array* res, const mlx_array a, int pad_width, const mlx_array pad_value, const char* mode, const mlx_stream s) {
-    return mlx_pad_symmetric_ptr(res, a, pad_width, pad_value, mode, s);
-}
-
-int mlx_partition_axis(mlx_array* res, const mlx_array a, int kth, int axis, const mlx_stream s) {
-    return mlx_partition_axis_ptr(res, a, kth, axis, s);
-}
-
-int mlx_partition(mlx_array* res, const mlx_array a, int kth, const mlx_stream s) {
-    return mlx_partition_ptr(res, a, kth, s);
-}
-
-int mlx_power(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_power_ptr(res, a, b, s);
-}
-
-int mlx_prod_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s) {
-    return mlx_prod_axes_ptr(res, a, axes, axes_num, keepdims, s);
-}
-
-int mlx_prod_axis(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s) {
-    return mlx_prod_axis_ptr(res, a, axis, keepdims, s);
-}
-
-int mlx_prod(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s) {
-    return mlx_prod_ptr(res, a, keepdims, s);
-}
-
-int mlx_put_along_axis(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array values, int axis, const mlx_stream s) {
-    return mlx_put_along_axis_ptr(res, a, indices, values, axis, s);
-}
-
-int mlx_qqmm(mlx_array* res, const mlx_array x, const mlx_array w, const mlx_array w_scales , mlx_optional_int group_size, mlx_optional_int bits, const char* mode, const mlx_stream s) {
-    return mlx_qqmm_ptr(res, x, w, w_scales, group_size, bits, mode, s);
-}
-
-int mlx_quantize(mlx_vector_array* res, const mlx_array w, mlx_optional_int group_size, mlx_optional_int bits, const char* mode, const mlx_stream s) {
-    return mlx_quantize_ptr(res, w, group_size, bits, mode, s);
-}
-
-int mlx_quantized_matmul(mlx_array* res, const mlx_array x, const mlx_array w, const mlx_array scales, const mlx_array biases , bool transpose, mlx_optional_int group_size, mlx_optional_int bits, const char* mode, const mlx_stream s) {
-    return mlx_quantized_matmul_ptr(res, x, w, scales, biases, transpose, group_size, bits, mode, s);
-}
-
-int mlx_radians(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_radians_ptr(res, a, s);
-}
-
-int mlx_real(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_real_ptr(res, a, s);
-}
-
-int mlx_reciprocal(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_reciprocal_ptr(res, a, s);
-}
-
-int mlx_remainder(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_remainder_ptr(res, a, b, s);
-}
-
-int mlx_repeat_axis(mlx_array* res, const mlx_array arr, int repeats, int axis, const mlx_stream s) {
-    return mlx_repeat_axis_ptr(res, arr, repeats, axis, s);
-}
-
-int mlx_repeat(mlx_array* res, const mlx_array arr, int repeats, const mlx_stream s) {
-    return mlx_repeat_ptr(res, arr, repeats, s);
-}
-
-int mlx_reshape(mlx_array* res, const mlx_array a, const int* shape, size_t shape_num, const mlx_stream s) {
-    return mlx_reshape_ptr(res, a, shape, shape_num, s);
-}
-
-int mlx_right_shift(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_right_shift_ptr(res, a, b, s);
-}
-
-int mlx_roll_axis(mlx_array* res, const mlx_array a, const int* shift, size_t shift_num, int axis, const mlx_stream s) {
-    return mlx_roll_axis_ptr(res, a, shift, shift_num, axis, s);
-}
-
-int mlx_roll_axes(mlx_array* res, const mlx_array a, const int* shift, size_t shift_num, const int* axes, size_t axes_num, const mlx_stream s) {
-    return mlx_roll_axes_ptr(res, a, shift, shift_num, axes, axes_num, s);
-}
-
-int mlx_roll(mlx_array* res, const mlx_array a, const int* shift, size_t shift_num, const mlx_stream s) {
-    return mlx_roll_ptr(res, a, shift, shift_num, s);
-}
-
-int mlx_round(mlx_array* res, const mlx_array a, int decimals, const mlx_stream s) {
-    return mlx_round_ptr(res, a, decimals, s);
-}
-
-int mlx_rsqrt(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_rsqrt_ptr(res, a, s);
-}
-
-int mlx_scatter(mlx_array* res, const mlx_array a, const mlx_vector_array indices, const mlx_array updates, const int* axes, size_t axes_num, const mlx_stream s) {
-    return mlx_scatter_ptr(res, a, indices, updates, axes, axes_num, s);
-}
-
-int mlx_scatter_single(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array updates, int axis, const mlx_stream s) {
-    return mlx_scatter_single_ptr(res, a, indices, updates, axis, s);
-}
-
-int mlx_scatter_add(mlx_array* res, const mlx_array a, const mlx_vector_array indices, const mlx_array updates, const int* axes, size_t axes_num, const mlx_stream s) {
-    return mlx_scatter_add_ptr(res, a, indices, updates, axes, axes_num, s);
-}
-
-int mlx_scatter_add_single(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array updates, int axis, const mlx_stream s) {
-    return mlx_scatter_add_single_ptr(res, a, indices, updates, axis, s);
-}
-
-int mlx_scatter_add_axis(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array values, int axis, const mlx_stream s) {
-    return mlx_scatter_add_axis_ptr(res, a, indices, values, axis, s);
-}
-
-int mlx_scatter_max(mlx_array* res, const mlx_array a, const mlx_vector_array indices, const mlx_array updates, const int* axes, size_t axes_num, const mlx_stream s) {
-    return mlx_scatter_max_ptr(res, a, indices, updates, axes, axes_num, s);
-}
-
-int mlx_scatter_max_single(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array updates, int axis, const mlx_stream s) {
-    return mlx_scatter_max_single_ptr(res, a, indices, updates, axis, s);
-}
-
-int mlx_scatter_min(mlx_array* res, const mlx_array a, const mlx_vector_array indices, const mlx_array updates, const int* axes, size_t axes_num, const mlx_stream s) {
-    return mlx_scatter_min_ptr(res, a, indices, updates, axes, axes_num, s);
-}
-
-int mlx_scatter_min_single(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array updates, int axis, const mlx_stream s) {
-    return mlx_scatter_min_single_ptr(res, a, indices, updates, axis, s);
-}
-
-int mlx_scatter_prod(mlx_array* res, const mlx_array a, const mlx_vector_array indices, const mlx_array updates, const int* axes, size_t axes_num, const mlx_stream s) {
-    return mlx_scatter_prod_ptr(res, a, indices, updates, axes, axes_num, s);
-}
-
-int mlx_scatter_prod_single(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array updates, int axis, const mlx_stream s) {
-    return mlx_scatter_prod_single_ptr(res, a, indices, updates, axis, s);
-}
-
-int mlx_segmented_mm(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_array segments, const mlx_stream s) {
-    return mlx_segmented_mm_ptr(res, a, b, segments, s);
-}
-
-int mlx_sigmoid(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_sigmoid_ptr(res, a, s);
-}
-
-int mlx_sign(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_sign_ptr(res, a, s);
-}
-
-int mlx_sin(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_sin_ptr(res, a, s);
-}
-
-int mlx_sinh(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_sinh_ptr(res, a, s);
-}
-
-int mlx_slice(mlx_array* res, const mlx_array a, const int* start, size_t start_num, const int* stop, size_t stop_num, const int* strides, size_t strides_num, const mlx_stream s) {
-    return mlx_slice_ptr(res, a, start, start_num, stop, stop_num, strides, strides_num, s);
-}
-
-int mlx_slice_dynamic(mlx_array* res, const mlx_array a, const mlx_array start, const int* axes, size_t axes_num, const int* slice_size, size_t slice_size_num, const mlx_stream s) {
-    return mlx_slice_dynamic_ptr(res, a, start, axes, axes_num, slice_size, slice_size_num, s);
-}
-
-int mlx_slice_update(mlx_array* res, const mlx_array src, const mlx_array update, const int* start, size_t start_num, const int* stop, size_t stop_num, const int* strides, size_t strides_num, const mlx_stream s) {
-    return mlx_slice_update_ptr(res, src, update, start, start_num, stop, stop_num, strides, strides_num, s);
-}
-
-int mlx_slice_update_dynamic(mlx_array* res, const mlx_array src, const mlx_array update, const mlx_array start, const int* axes, size_t axes_num, const mlx_stream s) {
-    return mlx_slice_update_dynamic_ptr(res, src, update, start, axes, axes_num, s);
-}
-
-int mlx_softmax_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool precise, const mlx_stream s) {
-    return mlx_softmax_axes_ptr(res, a, axes, axes_num, precise, s);
-}
-
-int mlx_softmax_axis(mlx_array* res, const mlx_array a, int axis, bool precise, const mlx_stream s) {
-    return mlx_softmax_axis_ptr(res, a, axis, precise, s);
-}
-
-int mlx_softmax(mlx_array* res, const mlx_array a, bool precise, const mlx_stream s) {
-    return mlx_softmax_ptr(res, a, precise, s);
-}
-
-int mlx_sort_axis(mlx_array* res, const mlx_array a, int axis, const mlx_stream s) {
-    return mlx_sort_axis_ptr(res, a, axis, s);
-}
-
-int mlx_sort(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_sort_ptr(res, a, s);
-}
-
-int mlx_split(mlx_vector_array* res, const mlx_array a, int num_splits, int axis, const mlx_stream s) {
-    return mlx_split_ptr(res, a, num_splits, axis, s);
-}
-
-int mlx_split_sections(mlx_vector_array* res, const mlx_array a, const int* indices, size_t indices_num, int axis, const mlx_stream s) {
-    return mlx_split_sections_ptr(res, a, indices, indices_num, axis, s);
-}
-
-int mlx_sqrt(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_sqrt_ptr(res, a, s);
-}
-
-int mlx_square(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_square_ptr(res, a, s);
-}
-
-int mlx_squeeze_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, const mlx_stream s) {
-    return mlx_squeeze_axes_ptr(res, a, axes, axes_num, s);
-}
-
-int mlx_squeeze_axis(mlx_array* res, const mlx_array a, int axis, const mlx_stream s) {
-    return mlx_squeeze_axis_ptr(res, a, axis, s);
-}
-
-int mlx_squeeze(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_squeeze_ptr(res, a, s);
-}
-
-int mlx_stack_axis(mlx_array* res, const mlx_vector_array arrays, int axis, const mlx_stream s) {
-    return mlx_stack_axis_ptr(res, arrays, axis, s);
-}
-
-int mlx_stack(mlx_array* res, const mlx_vector_array arrays, const mlx_stream s) {
-    return mlx_stack_ptr(res, arrays, s);
-}
-
-int mlx_std_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, int ddof, const mlx_stream s) {
-    return mlx_std_axes_ptr(res, a, axes, axes_num, keepdims, ddof, s);
-}
-
-int mlx_std_axis(mlx_array* res, const mlx_array a, int axis, bool keepdims, int ddof, const mlx_stream s) {
-    return mlx_std_axis_ptr(res, a, axis, keepdims, ddof, s);
-}
-
-int mlx_std(mlx_array* res, const mlx_array a, bool keepdims, int ddof, const mlx_stream s) {
-    return mlx_std_ptr(res, a, keepdims, ddof, s);
-}
-
-int mlx_stop_gradient(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_stop_gradient_ptr(res, a, s);
-}
-
-int mlx_subtract(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s) {
-    return mlx_subtract_ptr(res, a, b, s);
-}
-
-int mlx_sum_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s) {
-    return mlx_sum_axes_ptr(res, a, axes, axes_num, keepdims, s);
-}
-
-int mlx_sum_axis(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s) {
-    return mlx_sum_axis_ptr(res, a, axis, keepdims, s);
-}
-
-int mlx_sum(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s) {
-    return mlx_sum_ptr(res, a, keepdims, s);
-}
-
-int mlx_swapaxes(mlx_array* res, const mlx_array a, int axis1, int axis2, const mlx_stream s) {
-    return mlx_swapaxes_ptr(res, a, axis1, axis2, s);
-}
-
-int mlx_take_axis(mlx_array* res, const mlx_array a, const mlx_array indices, int axis, const mlx_stream s) {
-    return mlx_take_axis_ptr(res, a, indices, axis, s);
-}
-
-int mlx_take(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_stream s) {
-    return mlx_take_ptr(res, a, indices, s);
-}
-
-int mlx_take_along_axis(mlx_array* res, const mlx_array a, const mlx_array indices, int axis, const mlx_stream s) {
-    return mlx_take_along_axis_ptr(res, a, indices, axis, s);
-}
-
-int mlx_tan(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_tan_ptr(res, a, s);
-}
-
-int mlx_tanh(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_tanh_ptr(res, a, s);
-}
-
-int mlx_tensordot(mlx_array* res, const mlx_array a, const mlx_array b, const int* axes_a, size_t axes_a_num, const int* axes_b, size_t axes_b_num, const mlx_stream s) {
-    return mlx_tensordot_ptr(res, a, b, axes_a, axes_a_num, axes_b, axes_b_num, s);
-}
-
-int mlx_tensordot_axis(mlx_array* res, const mlx_array a, const mlx_array b, int axis, const mlx_stream s) {
-    return mlx_tensordot_axis_ptr(res, a, b, axis, s);
-}
-
-int mlx_tile(mlx_array* res, const mlx_array arr, const int* reps, size_t reps_num, const mlx_stream s) {
-    return mlx_tile_ptr(res, arr, reps, reps_num, s);
-}
-
-int mlx_to_fp8(mlx_array* res, const mlx_array x, const mlx_stream s) {
-    return mlx_to_fp8_ptr(res, x, s);
-}
-
-int mlx_topk_axis(mlx_array* res, const mlx_array a, int k, int axis, const mlx_stream s) {
-    return mlx_topk_axis_ptr(res, a, k, axis, s);
-}
-
-int mlx_topk(mlx_array* res, const mlx_array a, int k, const mlx_stream s) {
-    return mlx_topk_ptr(res, a, k, s);
-}
-
-int mlx_trace(mlx_array* res, const mlx_array a, int offset, int axis1, int axis2, mlx_dtype dtype, const mlx_stream s) {
-    return mlx_trace_ptr(res, a, offset, axis1, axis2, dtype, s);
-}
-
-int mlx_transpose_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, const mlx_stream s) {
-    return mlx_transpose_axes_ptr(res, a, axes, axes_num, s);
-}
-
-int mlx_transpose(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_transpose_ptr(res, a, s);
-}
-
-int mlx_tri(mlx_array* res, int n, int m, int k, mlx_dtype type, const mlx_stream s) {
-    return mlx_tri_ptr(res, n, m, k, type, s);
-}
-
-int mlx_tril(mlx_array* res, const mlx_array x, int k, const mlx_stream s) {
-    return mlx_tril_ptr(res, x, k, s);
-}
-
-int mlx_triu(mlx_array* res, const mlx_array x, int k, const mlx_stream s) {
-    return mlx_triu_ptr(res, x, k, s);
-}
-
-int mlx_unflatten(mlx_array* res, const mlx_array a, int axis, const int* shape, size_t shape_num, const mlx_stream s) {
-    return mlx_unflatten_ptr(res, a, axis, shape, shape_num, s);
-}
-
-int mlx_var_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, int ddof, const mlx_stream s) {
-    return mlx_var_axes_ptr(res, a, axes, axes_num, keepdims, ddof, s);
-}
-
-int mlx_var_axis(mlx_array* res, const mlx_array a, int axis, bool keepdims, int ddof, const mlx_stream s) {
-    return mlx_var_axis_ptr(res, a, axis, keepdims, ddof, s);
-}
-
-int mlx_var(mlx_array* res, const mlx_array a, bool keepdims, int ddof, const mlx_stream s) {
-    return mlx_var_ptr(res, a, keepdims, ddof, s);
-}
-
-int mlx_view(mlx_array* res, const mlx_array a, mlx_dtype dtype, const mlx_stream s) {
-    return mlx_view_ptr(res, a, dtype, s);
-}
-
-int mlx_where(mlx_array* res, const mlx_array condition, const mlx_array x, const mlx_array y, const mlx_stream s) {
-    return mlx_where_ptr(res, condition, x, y, s);
-}
-
-int mlx_zeros(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_stream s) {
-    return mlx_zeros_ptr(res, shape, shape_num, dtype, s);
-}
-
-int mlx_zeros_like(mlx_array* res, const mlx_array a, const mlx_stream s) {
-    return mlx_zeros_like_ptr(res, a, s);
-}
-
-int mlx_random_bernoulli(mlx_array* res, const mlx_array p, const int* shape, size_t shape_num, const mlx_array key , const mlx_stream s) {
-    return mlx_random_bernoulli_ptr(res, p, shape, shape_num, key, s);
-}
-
-int mlx_random_bits(mlx_array* res, const int* shape, size_t shape_num, int width, const mlx_array key , const mlx_stream s) {
-    return mlx_random_bits_ptr(res, shape, shape_num, width, key, s);
-}
-
-int mlx_random_categorical_shape(mlx_array* res, const mlx_array logits, int axis, const int* shape, size_t shape_num, const mlx_array key , const mlx_stream s) {
-    return mlx_random_categorical_shape_ptr(res, logits, axis, shape, shape_num, key, s);
-}
-
-int mlx_random_categorical_num_samples(mlx_array* res, const mlx_array logits_, int axis, int num_samples, const mlx_array key , const mlx_stream s) {
-    return mlx_random_categorical_num_samples_ptr(res, logits_, axis, num_samples, key, s);
-}
-
-int mlx_random_categorical(mlx_array* res, const mlx_array logits, int axis, const mlx_array key , const mlx_stream s) {
-    return mlx_random_categorical_ptr(res, logits, axis, key, s);
-}
-
-int mlx_random_gumbel(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_array key , const mlx_stream s) {
-    return mlx_random_gumbel_ptr(res, shape, shape_num, dtype, key, s);
-}
-
-int mlx_random_key(mlx_array* res, uint64_t seed) {
-    return mlx_random_key_ptr(res, seed);
-}
-
-int mlx_random_laplace(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, float loc, float scale, const mlx_array key , const mlx_stream s) {
-    return mlx_random_laplace_ptr(res, shape, shape_num, dtype, loc, scale, key, s);
-}
-
-int mlx_random_multivariate_normal(mlx_array* res, const mlx_array mean, const mlx_array cov, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_array key , const mlx_stream s) {
-    return mlx_random_multivariate_normal_ptr(res, mean, cov, shape, shape_num, dtype, key, s);
-}
-
-int mlx_random_normal_broadcast(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_array loc , const mlx_array scale , const mlx_array key , const mlx_stream s) {
-    return mlx_random_normal_broadcast_ptr(res, shape, shape_num, dtype, loc, scale, key, s);
-}
-
-int mlx_random_normal(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, float loc, float scale, const mlx_array key , const mlx_stream s) {
-    return mlx_random_normal_ptr(res, shape, shape_num, dtype, loc, scale, key, s);
-}
-
-int mlx_random_permutation(mlx_array* res, const mlx_array x, int axis, const mlx_array key , const mlx_stream s) {
-    return mlx_random_permutation_ptr(res, x, axis, key, s);
-}
-
-int mlx_random_permutation_arange(mlx_array* res, int x, const mlx_array key , const mlx_stream s) {
-    return mlx_random_permutation_arange_ptr(res, x, key, s);
-}
-
-int mlx_random_randint(mlx_array* res, const mlx_array low, const mlx_array high, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_array key , const mlx_stream s) {
-    return mlx_random_randint_ptr(res, low, high, shape, shape_num, dtype, key, s);
-}
-
-int mlx_random_seed(uint64_t seed) {
-    return mlx_random_seed_ptr(seed);
-}
-
-int mlx_random_split_num(mlx_array* res, const mlx_array key, int num, const mlx_stream s) {
-    return mlx_random_split_num_ptr(res, key, num, s);
-}
-
-int mlx_random_split(mlx_array* res_0, mlx_array* res_1, const mlx_array key, const mlx_stream s) {
-    return mlx_random_split_ptr(res_0, res_1, key, s);
-}
-
-int mlx_random_truncated_normal(mlx_array* res, const mlx_array lower, const mlx_array upper, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_array key , const mlx_stream s) {
-    return mlx_random_truncated_normal_ptr(res, lower, upper, shape, shape_num, dtype, key, s);
-}
-
-int mlx_random_uniform(mlx_array* res, const mlx_array low, const mlx_array high, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_array key , const mlx_stream s) {
-    return mlx_random_uniform_ptr(res, low, high, shape, shape_num, dtype, key, s);
-}
-
-mlx_stream mlx_stream_new(void) {
-    return mlx_stream_new_ptr();
-}
-
-mlx_stream mlx_stream_new_device(mlx_device dev) {
-    return mlx_stream_new_device_ptr(dev);
-}
-
-int mlx_stream_set(mlx_stream* stream, const mlx_stream src) {
-    return mlx_stream_set_ptr(stream, src);
-}
-
-int mlx_stream_free(mlx_stream stream) {
-    return mlx_stream_free_ptr(stream);
-}
-
-int mlx_stream_tostring(mlx_string* str, mlx_stream stream) {
-    return mlx_stream_tostring_ptr(str, stream);
-}
-
-bool mlx_stream_equal(mlx_stream lhs, mlx_stream rhs) {
-    return mlx_stream_equal_ptr(lhs, rhs);
-}
-
-int mlx_stream_get_device(mlx_device* dev, mlx_stream stream) {
-    return mlx_stream_get_device_ptr(dev, stream);
-}
-
-int mlx_stream_get_index(int* index, mlx_stream stream) {
-    return mlx_stream_get_index_ptr(index, stream);
-}
-
-int mlx_synchronize(mlx_stream stream) {
-    return mlx_synchronize_ptr(stream);
-}
-
-int mlx_get_default_stream(mlx_stream* stream, mlx_device dev) {
-    return mlx_get_default_stream_ptr(stream, dev);
-}
-
-int mlx_set_default_stream(mlx_stream stream) {
-    return mlx_set_default_stream_ptr(stream);
-}
-
-mlx_stream mlx_default_cpu_stream_new(void) {
-    return mlx_default_cpu_stream_new_ptr();
-}
-
-mlx_stream mlx_default_gpu_stream_new(void) {
-    return mlx_default_gpu_stream_new_ptr();
-}
-
-mlx_string mlx_string_new(void) {
-    return mlx_string_new_ptr();
-}
-
-mlx_string mlx_string_new_data(const char* str) {
-    return mlx_string_new_data_ptr(str);
-}
-
-int mlx_string_set(mlx_string* str, const mlx_string src) {
-    return mlx_string_set_ptr(str, src);
-}
-
-const char* mlx_string_data(mlx_string str) {
-    return mlx_string_data_ptr(str);
-}
-
-int mlx_string_free(mlx_string str) {
-    return mlx_string_free_ptr(str);
-}
-
-int mlx_async_eval(const mlx_vector_array outputs) {
-    return mlx_async_eval_ptr(outputs);
-}
-
-int mlx_checkpoint(mlx_closure* res, const mlx_closure fun) {
-    return mlx_checkpoint_ptr(res, fun);
-}
-
-int mlx_custom_function(mlx_closure* res, const mlx_closure fun, const mlx_closure_custom fun_vjp , const mlx_closure_custom_jvp fun_jvp , const mlx_closure_custom_vmap fun_vmap) {
-    return mlx_custom_function_ptr(res, fun, fun_vjp, fun_jvp, fun_vmap);
-}
-
-int mlx_custom_vjp(mlx_closure* res, const mlx_closure fun, const mlx_closure_custom fun_vjp) {
-    return mlx_custom_vjp_ptr(res, fun, fun_vjp);
-}
-
-int mlx_eval(const mlx_vector_array outputs) {
-    return mlx_eval_ptr(outputs);
-}
-
-int mlx_jvp(mlx_vector_array* res_0, mlx_vector_array* res_1, const mlx_closure fun, const mlx_vector_array primals, const mlx_vector_array tangents) {
-    return mlx_jvp_ptr(res_0, res_1, fun, primals, tangents);
-}
-
-int mlx_value_and_grad(mlx_closure_value_and_grad* res, const mlx_closure fun, const int* argnums, size_t argnums_num) {
-    return mlx_value_and_grad_ptr(res, fun, argnums, argnums_num);
-}
-
-int mlx_vjp(mlx_vector_array* res_0, mlx_vector_array* res_1, const mlx_closure fun, const mlx_vector_array primals, const mlx_vector_array cotangents) {
-    return mlx_vjp_ptr(res_0, res_1, fun, primals, cotangents);
-}
-
-int mlx_detail_vmap_replace(mlx_vector_array* res, const mlx_vector_array inputs, const mlx_vector_array s_inputs, const mlx_vector_array s_outputs, const int* in_axes, size_t in_axes_num, const int* out_axes, size_t out_axes_num) {
-    return mlx_detail_vmap_replace_ptr(res, inputs, s_inputs, s_outputs, in_axes, in_axes_num, out_axes, out_axes_num);
-}
-
-int mlx_detail_vmap_trace(mlx_vector_array* res_0, mlx_vector_array* res_1, const mlx_closure fun, const mlx_vector_array inputs, const int* in_axes, size_t in_axes_num) {
-    return mlx_detail_vmap_trace_ptr(res_0, res_1, fun, inputs, in_axes, in_axes_num);
-}
-
-mlx_vector_array mlx_vector_array_new(void) {
-    return mlx_vector_array_new_ptr();
-}
-
-int mlx_vector_array_set(mlx_vector_array* vec, const mlx_vector_array src) {
-    return mlx_vector_array_set_ptr(vec, src);
-}
-
-int mlx_vector_array_free(mlx_vector_array vec) {
-    return mlx_vector_array_free_ptr(vec);
-}
-
-mlx_vector_array mlx_vector_array_new_data(const mlx_array* data, size_t size) {
-    return mlx_vector_array_new_data_ptr(data, size);
-}
-
-mlx_vector_array mlx_vector_array_new_value(const mlx_array val) {
-    return mlx_vector_array_new_value_ptr(val);
-}
-
-int mlx_vector_array_set_data(mlx_vector_array* vec, const mlx_array* data, size_t size) {
-    return mlx_vector_array_set_data_ptr(vec, data, size);
-}
-
-int mlx_vector_array_set_value(mlx_vector_array* vec, const mlx_array val) {
-    return mlx_vector_array_set_value_ptr(vec, val);
-}
-
-int mlx_vector_array_append_data(mlx_vector_array vec, const mlx_array* data, size_t size) {
-    return mlx_vector_array_append_data_ptr(vec, data, size);
-}
-
-int mlx_vector_array_append_value(mlx_vector_array vec, const mlx_array val) {
-    return mlx_vector_array_append_value_ptr(vec, val);
-}
-
-size_t mlx_vector_array_size(mlx_vector_array vec) {
-    return mlx_vector_array_size_ptr(vec);
-}
-
-int mlx_vector_array_get(mlx_array* res, const mlx_vector_array vec, size_t idx) {
-    return mlx_vector_array_get_ptr(res, vec, idx);
-}
-
-mlx_vector_vector_array mlx_vector_vector_array_new(void) {
-    return mlx_vector_vector_array_new_ptr();
-}
-
-int mlx_vector_vector_array_set(mlx_vector_vector_array* vec, const mlx_vector_vector_array src) {
-    return mlx_vector_vector_array_set_ptr(vec, src);
-}
-
-int mlx_vector_vector_array_free(mlx_vector_vector_array vec) {
-    return mlx_vector_vector_array_free_ptr(vec);
-}
-
-mlx_vector_vector_array mlx_vector_vector_array_new_data(const mlx_vector_array* data, size_t size) {
-    return mlx_vector_vector_array_new_data_ptr(data, size);
-}
-
-mlx_vector_vector_array mlx_vector_vector_array_new_value(const mlx_vector_array val) {
-    return mlx_vector_vector_array_new_value_ptr(val);
-}
-
-int mlx_vector_vector_array_set_data(mlx_vector_vector_array* vec, const mlx_vector_array* data, size_t size) {
-    return mlx_vector_vector_array_set_data_ptr(vec, data, size);
-}
-
-int mlx_vector_vector_array_set_value(mlx_vector_vector_array* vec, const mlx_vector_array val) {
-    return mlx_vector_vector_array_set_value_ptr(vec, val);
-}
-
-int mlx_vector_vector_array_append_data(mlx_vector_vector_array vec, const mlx_vector_array* data, size_t size) {
-    return mlx_vector_vector_array_append_data_ptr(vec, data, size);
-}
-
-int mlx_vector_vector_array_append_value(mlx_vector_vector_array vec, const mlx_vector_array val) {
-    return mlx_vector_vector_array_append_value_ptr(vec, val);
-}
-
-size_t mlx_vector_vector_array_size(mlx_vector_vector_array vec) {
-    return mlx_vector_vector_array_size_ptr(vec);
-}
-
-int mlx_vector_vector_array_get(mlx_vector_array* res, const mlx_vector_vector_array vec, size_t idx) {
-    return mlx_vector_vector_array_get_ptr(res, vec, idx);
-}
-
-mlx_vector_int mlx_vector_int_new(void) {
-    return mlx_vector_int_new_ptr();
-}
-
-int mlx_vector_int_set(mlx_vector_int* vec, const mlx_vector_int src) {
-    return mlx_vector_int_set_ptr(vec, src);
-}
-
-int mlx_vector_int_free(mlx_vector_int vec) {
-    return mlx_vector_int_free_ptr(vec);
-}
-
-mlx_vector_int mlx_vector_int_new_data(int* data, size_t size) {
-    return mlx_vector_int_new_data_ptr(data, size);
-}
-
-mlx_vector_int mlx_vector_int_new_value(int val) {
-    return mlx_vector_int_new_value_ptr(val);
-}
-
-int mlx_vector_int_set_data(mlx_vector_int* vec, int* data, size_t size) {
-    return mlx_vector_int_set_data_ptr(vec, data, size);
-}
-
-int mlx_vector_int_set_value(mlx_vector_int* vec, int val) {
-    return mlx_vector_int_set_value_ptr(vec, val);
-}
-
-int mlx_vector_int_append_data(mlx_vector_int vec, int* data, size_t size) {
-    return mlx_vector_int_append_data_ptr(vec, data, size);
-}
-
-int mlx_vector_int_append_value(mlx_vector_int vec, int val) {
-    return mlx_vector_int_append_value_ptr(vec, val);
-}
-
-size_t mlx_vector_int_size(mlx_vector_int vec) {
-    return mlx_vector_int_size_ptr(vec);
-}
-
-int mlx_vector_int_get(int* res, const mlx_vector_int vec, size_t idx) {
-    return mlx_vector_int_get_ptr(res, vec, idx);
-}
-
-mlx_vector_string mlx_vector_string_new(void) {
-    return mlx_vector_string_new_ptr();
-}
-
-int mlx_vector_string_set(mlx_vector_string* vec, const mlx_vector_string src) {
-    return mlx_vector_string_set_ptr(vec, src);
-}
-
-int mlx_vector_string_free(mlx_vector_string vec) {
-    return mlx_vector_string_free_ptr(vec);
-}
-
-mlx_vector_string mlx_vector_string_new_data(const char** data, size_t size) {
-    return mlx_vector_string_new_data_ptr(data, size);
-}
-
-mlx_vector_string mlx_vector_string_new_value(const char* val) {
-    return mlx_vector_string_new_value_ptr(val);
-}
-
-int mlx_vector_string_set_data(mlx_vector_string* vec, const char** data, size_t size) {
-    return mlx_vector_string_set_data_ptr(vec, data, size);
-}
-
-int mlx_vector_string_set_value(mlx_vector_string* vec, const char* val) {
-    return mlx_vector_string_set_value_ptr(vec, val);
-}
-
-int mlx_vector_string_append_data(mlx_vector_string vec, const char** data, size_t size) {
-    return mlx_vector_string_append_data_ptr(vec, data, size);
-}
-
-int mlx_vector_string_append_value(mlx_vector_string vec, const char* val) {
-    return mlx_vector_string_append_value_ptr(vec, val);
-}
-
-size_t mlx_vector_string_size(mlx_vector_string vec) {
-    return mlx_vector_string_size_ptr(vec);
-}
-
-int mlx_vector_string_get(char** res, const mlx_vector_string vec, size_t idx) {
-    return mlx_vector_string_get_ptr(res, vec, idx);
-}
-
-int mlx_version(mlx_string* str_) {
-    return mlx_version_ptr(str_);
-}
-

x/imagegen/mlx/mlx.go

@@ -3,13 +3,12 @@
 package mlx
 
 /*
-#cgo CFLAGS: -O3 -I${SRCDIR}/../../../build/_deps/mlx-c-src -I${SRCDIR}
+#cgo CFLAGS: -O3 -I${SRCDIR}/../../../build/_deps/mlx-c-src
+#cgo LDFLAGS: -L${SRCDIR}/../../../build/lib/ollama/ -lmlxc -Wl,-rpath,${SRCDIR}/../../../build/lib/ollama/
 #cgo darwin LDFLAGS: -lc++ -framework Metal -framework Foundation -framework Accelerate
-#cgo linux LDFLAGS: -lstdc++ -ldl
-#cgo windows LDFLAGS: -lstdc++
+#cgo linux LDFLAGS: -lstdc++ -lcuda -lcudart -lnvrtc
 
-// Use generated wrappers instead of direct MLX headers
-#include "mlx.h"
+#include "mlx/c/mlx.h"
 #include <stdlib.h>
 #include <stdint.h>
 #include <string.h>
@@ -43,6 +42,192 @@ static inline mlx_stream cpu_stream() {
 // CGO noescape/nocallback hints to reduce CGO overhead
 // noescape: pointers won't escape, no heap allocation needed
 // nocallback: function won't call back into Go
+#cgo noescape mlx_add
+#cgo nocallback mlx_add
+#cgo noescape mlx_subtract
+#cgo nocallback mlx_subtract
+#cgo noescape mlx_multiply
+#cgo nocallback mlx_multiply
+#cgo noescape mlx_divide
+#cgo nocallback mlx_divide
+#cgo noescape mlx_negative
+#cgo nocallback mlx_negative
+#cgo noescape mlx_abs
+#cgo nocallback mlx_abs
+#cgo noescape mlx_exp
+#cgo nocallback mlx_exp
+#cgo noescape mlx_log
+#cgo nocallback mlx_log
+#cgo noescape mlx_sqrt
+#cgo nocallback mlx_sqrt
+#cgo noescape mlx_rsqrt
+#cgo nocallback mlx_rsqrt
+#cgo noescape mlx_square
+#cgo nocallback mlx_square
+#cgo noescape mlx_power
+#cgo nocallback mlx_power
+#cgo noescape mlx_erf
+#cgo nocallback mlx_erf
+#cgo noescape mlx_sigmoid
+#cgo nocallback mlx_sigmoid
+#cgo noescape mlx_tanh
+#cgo nocallback mlx_tanh
+#cgo noescape mlx_sin
+#cgo nocallback mlx_sin
+#cgo noescape mlx_cos
+#cgo nocallback mlx_cos
+#cgo noescape mlx_maximum
+#cgo nocallback mlx_maximum
+#cgo noescape mlx_minimum
+#cgo nocallback mlx_minimum
+#cgo noescape mlx_clip
+#cgo nocallback mlx_clip
+#cgo noescape mlx_sum
+#cgo nocallback mlx_sum
+#cgo noescape mlx_sum_axis
+#cgo nocallback mlx_sum_axis
+#cgo noescape mlx_mean
+#cgo nocallback mlx_mean
+#cgo noescape mlx_mean_axis
+#cgo nocallback mlx_mean_axis
+#cgo noescape mlx_var_axis
+#cgo nocallback mlx_var_axis
+#cgo noescape mlx_argmax
+#cgo nocallback mlx_argmax
+#cgo noescape mlx_argmax_axis
+#cgo nocallback mlx_argmax_axis
+#cgo noescape mlx_softmax_axis
+#cgo nocallback mlx_softmax_axis
+#cgo noescape mlx_cumsum
+#cgo nocallback mlx_cumsum
+#cgo noescape mlx_matmul
+#cgo nocallback mlx_matmul
+#cgo noescape mlx_addmm
+#cgo nocallback mlx_addmm
+#cgo noescape mlx_gather_mm
+#cgo nocallback mlx_gather_mm
+#cgo noescape mlx_gather_qmm
+#cgo nocallback mlx_gather_qmm
+#cgo noescape mlx_reshape
+#cgo nocallback mlx_reshape
+#cgo noescape mlx_transpose_axes
+#cgo nocallback mlx_transpose_axes
+#cgo noescape mlx_expand_dims
+#cgo nocallback mlx_expand_dims
+#cgo noescape mlx_squeeze_axis
+#cgo nocallback mlx_squeeze_axis
+#cgo noescape mlx_flatten
+#cgo nocallback mlx_flatten
+#cgo noescape mlx_concatenate_axis
+#cgo nocallback mlx_concatenate_axis
+#cgo noescape mlx_slice
+#cgo nocallback mlx_slice
+#cgo noescape mlx_slice_update
+#cgo nocallback mlx_slice_update
+#cgo noescape mlx_as_strided
+#cgo nocallback mlx_as_strided
+#cgo noescape mlx_view
+#cgo nocallback mlx_view
+#cgo noescape mlx_contiguous
+#cgo nocallback mlx_contiguous
+#cgo noescape mlx_pad
+#cgo nocallback mlx_pad
+#cgo noescape mlx_tile
+#cgo nocallback mlx_tile
+#cgo noescape mlx_take_axis
+#cgo nocallback mlx_take_axis
+#cgo noescape mlx_take_along_axis
+#cgo nocallback mlx_take_along_axis
+#cgo noescape mlx_put_along_axis
+#cgo nocallback mlx_put_along_axis
+#cgo noescape mlx_where
+#cgo nocallback mlx_where
+#cgo noescape mlx_argsort_axis
+#cgo nocallback mlx_argsort_axis
+#cgo noescape mlx_argpartition_axis
+#cgo nocallback mlx_argpartition_axis
+#cgo noescape mlx_topk_axis
+#cgo nocallback mlx_topk_axis
+#cgo noescape mlx_less
+#cgo nocallback mlx_less
+#cgo noescape mlx_greater_equal
+#cgo nocallback mlx_greater_equal
+#cgo noescape mlx_logical_and
+#cgo nocallback mlx_logical_and
+#cgo noescape mlx_zeros
+#cgo nocallback mlx_zeros
+#cgo noescape mlx_zeros_like
+#cgo nocallback mlx_zeros_like
+#cgo noescape mlx_ones
+#cgo nocallback mlx_ones
+#cgo noescape mlx_full
+#cgo nocallback mlx_full
+#cgo noescape mlx_arange
+#cgo nocallback mlx_arange
+#cgo noescape mlx_linspace
+#cgo nocallback mlx_linspace
+#cgo noescape mlx_tri
+#cgo nocallback mlx_tri
+#cgo noescape mlx_astype
+#cgo nocallback mlx_astype
+#cgo noescape mlx_fast_rms_norm
+#cgo nocallback mlx_fast_rms_norm
+#cgo noescape mlx_fast_rope
+#cgo nocallback mlx_fast_rope
+#cgo noescape mlx_fast_scaled_dot_product_attention
+#cgo nocallback mlx_fast_scaled_dot_product_attention
+#cgo noescape mlx_conv2d
+#cgo nocallback mlx_conv2d
+#cgo noescape mlx_conv3d
+#cgo nocallback mlx_conv3d
+#cgo noescape mlx_random_key
+#cgo nocallback mlx_random_key
+#cgo noescape mlx_random_split
+#cgo nocallback mlx_random_split
+#cgo noescape mlx_random_categorical_num_samples
+#cgo nocallback mlx_random_categorical_num_samples
+#cgo noescape mlx_random_normal
+#cgo nocallback mlx_random_normal
+#cgo noescape mlx_random_uniform
+#cgo nocallback mlx_random_uniform
+#cgo noescape mlx_array_eval
+#cgo nocallback mlx_array_eval
+#cgo noescape mlx_eval
+#cgo nocallback mlx_eval
+#cgo noescape mlx_async_eval
+#cgo nocallback mlx_async_eval
+#cgo noescape mlx_synchronize
+#cgo nocallback mlx_synchronize
+#cgo noescape mlx_array_new
+#cgo nocallback mlx_array_new
+#cgo noescape mlx_array_new_data
+#cgo nocallback mlx_array_new_data
+#cgo noescape mlx_array_new_float
+#cgo nocallback mlx_array_new_float
+#cgo noescape mlx_array_free
+#cgo nocallback mlx_array_free
+#cgo noescape mlx_array_size
+#cgo nocallback mlx_array_size
+#cgo noescape mlx_array_ndim
+#cgo nocallback mlx_array_ndim
+#cgo noescape mlx_array_dim
+#cgo nocallback mlx_array_dim
+#cgo noescape mlx_array_dtype
+#cgo nocallback mlx_array_dtype
+#cgo noescape mlx_array_item_int32
+#cgo nocallback mlx_array_item_int32
+#cgo noescape mlx_vector_array_new_data
+#cgo nocallback mlx_vector_array_new_data
+#cgo noescape mlx_vector_array_free
+#cgo nocallback mlx_vector_array_free
+#cgo noescape mlx_array_new_int
+#cgo nocallback mlx_array_new_int
+#cgo noescape mlx_stream_new_device
+#cgo nocallback mlx_stream_new_device
+#cgo noescape mlx_get_default_stream
+#cgo nocallback mlx_get_default_stream
+#cgo noescape mlx_set_default_stream
+#cgo nocallback mlx_set_default_stream
 */
 import "C"
 import (
@@ -1137,27 +1322,6 @@ func RMSNormNoWeight(x *Array, eps float32) *Array {
 	return RMSNorm(x, ones, eps)
 }
 
-// LayerNorm applies layer normalization without learnable params
-// (x - mean) / sqrt(var + eps)
-func LayerNorm(x *Array, eps float32) *Array {
-	return LayerNormWithWeightBias(x, nil, nil, eps)
-}
-
-// LayerNormWithWeightBias computes layer normalization using mlx.fast
-// weight and bias can be nil for elementwise_affine=False
-func LayerNormWithWeightBias(x, weight, bias *Array, eps float32) *Array {
-	res := C.mlx_array_new()
-	var wc, bc C.mlx_array
-	if weight != nil {
-		wc = weight.c
-	}
-	if bias != nil {
-		bc = bias.c
-	}
-	C.mlx_fast_layer_norm(&res, x.c, wc, bc, C.float(eps), C.default_stream())
-	return newArray(res)
-}
-
 // RoPE applies rotary position embeddings using mlx.fast
 func RoPE(x *Array, dims int, traditional bool, base, scale float32, offset int) *Array {
 	res := C.mlx_array_new()
@@ -1632,57 +1796,7 @@ func ArgmaxKeepArray(logits *Array) *Array {
 var RandomState = []*Array{nil}
 var randomStateMu sync.Mutex
 
-var mlxInitialized bool
-var mlxInitError error
-
-// InitMLX initializes the MLX library by dynamically loading libmlxc.
-// This must be called before using any MLX functions.
-// Returns an error if the library cannot be loaded.
-func InitMLX() error {
-	if mlxInitialized {
-		return mlxInitError
-	}
-
-	// Try to load the MLX dynamic library
-	ret := C.mlx_dynamic_init()
-	if ret != 0 {
-		errMsg := C.GoString(C.mlx_dynamic_error())
-		mlxInitError = fmt.Errorf("failed to initialize MLX: %s", errMsg)
-		return mlxInitError
-	}
-
-	// Initialize all function pointers via dlsym
-	handle := C.mlx_get_handle()
-	ret = C.mlx_load_functions(handle)
-	if ret != 0 {
-		mlxInitError = fmt.Errorf("failed to load MLX function symbols")
-		return mlxInitError
-	}
-
-	mlxInitialized = true
-	mlxInitError = nil
-	return nil
-}
-
-// IsMLXAvailable returns whether MLX was successfully initialized
-func IsMLXAvailable() bool {
-	return mlxInitialized && mlxInitError == nil
-}
-
-// GetMLXInitError returns any error that occurred during MLX initialization
-func GetMLXInitError() error {
-	return mlxInitError
-}
-
 func init() {
-	// Initialize MLX dynamic library first
-	if err := InitMLX(); err != nil {
-		// Don't panic in init - let the caller handle the error
-		// Store the error for later retrieval
-		mlxInitError = err
-		return
-	}
-
 	// Lock main goroutine to OS thread for CUDA context stability.
 	// CUDA contexts are bound to threads; Go can migrate goroutines between threads.
 	runtime.LockOSThread()

x/imagegen/mlx/mlx.h

@@ -1,2337 +0,0 @@
-// AUTO-GENERATED by generate_wrappers.go - DO NOT EDIT
-// This file provides wrapper declarations for MLX-C functions that use dlopen/dlsym
-//
-// Strategy: Include MLX-C headers for type definitions, then provide wrapper
-// functions that shadow the originals, allowing Go code to call them directly (e.g., C.mlx_add).
-// Function pointers are defined in mlx.c (single compilation unit).
-
-#ifndef MLX_WRAPPERS_H
-#define MLX_WRAPPERS_H
-
-// Include MLX headers for type definitions and original declarations
-#include "mlx/c/mlx.h"
-#include "mlx_dynamic.h"
-#include <stdio.h>
-
-// Undefine any existing MLX function macros
-#undef mlx_dtype_size
-#undef mlx_array_tostring
-#undef mlx_array_new
-#undef mlx_array_free
-#undef mlx_array_new_bool
-#undef mlx_array_new_int
-#undef mlx_array_new_float32
-#undef mlx_array_new_float
-#undef mlx_array_new_float64
-#undef mlx_array_new_double
-#undef mlx_array_new_complex
-#undef mlx_array_new_data
-#undef mlx_array_set
-#undef mlx_array_set_bool
-#undef mlx_array_set_int
-#undef mlx_array_set_float32
-#undef mlx_array_set_float
-#undef mlx_array_set_float64
-#undef mlx_array_set_double
-#undef mlx_array_set_complex
-#undef mlx_array_set_data
-#undef mlx_array_itemsize
-#undef mlx_array_size
-#undef mlx_array_nbytes
-#undef mlx_array_ndim
-#undef mlx_array_shape
-#undef mlx_array_strides
-#undef mlx_array_dim
-#undef mlx_array_dtype
-#undef mlx_array_eval
-#undef mlx_array_item_bool
-#undef mlx_array_item_uint8
-#undef mlx_array_item_uint16
-#undef mlx_array_item_uint32
-#undef mlx_array_item_uint64
-#undef mlx_array_item_int8
-#undef mlx_array_item_int16
-#undef mlx_array_item_int32
-#undef mlx_array_item_int64
-#undef mlx_array_item_float32
-#undef mlx_array_item_float64
-#undef mlx_array_item_complex64
-#undef mlx_array_item_float16
-#undef mlx_array_item_bfloat16
-#undef mlx_array_data_bool
-#undef mlx_array_data_uint8
-#undef mlx_array_data_uint16
-#undef mlx_array_data_uint32
-#undef mlx_array_data_uint64
-#undef mlx_array_data_int8
-#undef mlx_array_data_int16
-#undef mlx_array_data_int32
-#undef mlx_array_data_int64
-#undef mlx_array_data_float32
-#undef mlx_array_data_float64
-#undef mlx_array_data_complex64
-#undef mlx_array_data_float16
-#undef mlx_array_data_bfloat16
-#undef _mlx_array_is_available
-#undef _mlx_array_wait
-#undef _mlx_array_is_contiguous
-#undef _mlx_array_is_row_contiguous
-#undef _mlx_array_is_col_contiguous
-#undef mlx_closure_new
-#undef mlx_closure_free
-#undef mlx_closure_new_func
-#undef mlx_closure_new_func_payload
-#undef mlx_closure_set
-#undef mlx_closure_apply
-#undef mlx_closure_new_unary
-#undef mlx_closure_kwargs_new
-#undef mlx_closure_kwargs_free
-#undef mlx_closure_kwargs_new_func
-#undef mlx_closure_kwargs_new_func_payload
-#undef mlx_closure_kwargs_set
-#undef mlx_closure_kwargs_apply
-#undef mlx_closure_value_and_grad_new
-#undef mlx_closure_value_and_grad_free
-#undef mlx_closure_value_and_grad_new_func
-#undef mlx_closure_value_and_grad_new_func_payload
-#undef mlx_closure_value_and_grad_set
-#undef mlx_closure_value_and_grad_apply
-#undef mlx_closure_custom_new
-#undef mlx_closure_custom_free
-#undef mlx_closure_custom_new_func
-#undef mlx_closure_custom_new_func_payload
-#undef mlx_closure_custom_set
-#undef mlx_closure_custom_apply
-#undef mlx_closure_custom_jvp_new
-#undef mlx_closure_custom_jvp_free
-#undef mlx_closure_custom_jvp_new_func
-#undef mlx_closure_custom_jvp_new_func_payload
-#undef mlx_closure_custom_jvp_set
-#undef mlx_closure_custom_jvp_apply
-#undef mlx_closure_custom_vmap_new
-#undef mlx_closure_custom_vmap_free
-#undef mlx_closure_custom_vmap_new_func
-#undef mlx_closure_custom_vmap_new_func_payload
-#undef mlx_closure_custom_vmap_set
-#undef mlx_closure_custom_vmap_apply
-#undef mlx_compile
-#undef mlx_detail_compile
-#undef mlx_detail_compile_clear_cache
-#undef mlx_detail_compile_erase
-#undef mlx_disable_compile
-#undef mlx_enable_compile
-#undef mlx_set_compile_mode
-#undef mlx_device_new
-#undef mlx_device_new_type
-#undef mlx_device_free
-#undef mlx_device_set
-#undef mlx_device_tostring
-#undef mlx_device_equal
-#undef mlx_device_get_index
-#undef mlx_device_get_type
-#undef mlx_get_default_device
-#undef mlx_set_default_device
-#undef mlx_distributed_all_gather
-#undef mlx_distributed_all_max
-#undef mlx_distributed_all_min
-#undef mlx_distributed_all_sum
-#undef mlx_distributed_recv
-#undef mlx_distributed_recv_like
-#undef mlx_distributed_send
-#undef mlx_distributed_sum_scatter
-#undef mlx_distributed_group_rank
-#undef mlx_distributed_group_size
-#undef mlx_distributed_group_split
-#undef mlx_distributed_is_available
-#undef mlx_distributed_init
-#undef mlx_set_error_handler
-#undef _mlx_error
-#undef mlx_export_function
-#undef mlx_export_function_kwargs
-#undef mlx_function_exporter_new
-#undef mlx_function_exporter_free
-#undef mlx_function_exporter_apply
-#undef mlx_function_exporter_apply_kwargs
-#undef mlx_imported_function_new
-#undef mlx_imported_function_free
-#undef mlx_imported_function_apply
-#undef mlx_imported_function_apply_kwargs
-#undef mlx_fast_cuda_kernel_config_new
-#undef mlx_fast_cuda_kernel_config_free
-#undef mlx_fast_cuda_kernel_config_add_output_arg
-#undef mlx_fast_cuda_kernel_config_set_grid
-#undef mlx_fast_cuda_kernel_config_set_thread_group
-#undef mlx_fast_cuda_kernel_config_set_init_value
-#undef mlx_fast_cuda_kernel_config_set_verbose
-#undef mlx_fast_cuda_kernel_config_add_template_arg_dtype
-#undef mlx_fast_cuda_kernel_config_add_template_arg_int
-#undef mlx_fast_cuda_kernel_config_add_template_arg_bool
-#undef mlx_fast_cuda_kernel_new
-#undef mlx_fast_cuda_kernel_free
-#undef mlx_fast_cuda_kernel_apply
-#undef mlx_fast_layer_norm
-#undef mlx_fast_metal_kernel_config_new
-#undef mlx_fast_metal_kernel_config_free
-#undef mlx_fast_metal_kernel_config_add_output_arg
-#undef mlx_fast_metal_kernel_config_set_grid
-#undef mlx_fast_metal_kernel_config_set_thread_group
-#undef mlx_fast_metal_kernel_config_set_init_value
-#undef mlx_fast_metal_kernel_config_set_verbose
-#undef mlx_fast_metal_kernel_config_add_template_arg_dtype
-#undef mlx_fast_metal_kernel_config_add_template_arg_int
-#undef mlx_fast_metal_kernel_config_add_template_arg_bool
-#undef mlx_fast_metal_kernel_new
-#undef mlx_fast_metal_kernel_free
-#undef mlx_fast_metal_kernel_apply
-#undef mlx_fast_rms_norm
-#undef mlx_fast_rope
-#undef mlx_fast_rope_dynamic
-#undef mlx_fast_scaled_dot_product_attention
-#undef mlx_fft_fft
-#undef mlx_fft_fft2
-#undef mlx_fft_fftn
-#undef mlx_fft_fftshift
-#undef mlx_fft_ifft
-#undef mlx_fft_ifft2
-#undef mlx_fft_ifftn
-#undef mlx_fft_ifftshift
-#undef mlx_fft_irfft
-#undef mlx_fft_irfft2
-#undef mlx_fft_irfftn
-#undef mlx_fft_rfft
-#undef mlx_fft_rfft2
-#undef mlx_fft_rfftn
-#undef mlx_load_reader
-#undef mlx_load
-#undef mlx_load_safetensors_reader
-#undef mlx_load_safetensors
-#undef mlx_save_writer
-#undef mlx_save
-#undef mlx_save_safetensors_writer
-#undef mlx_save_safetensors
-#undef mlx_io_reader_new
-#undef mlx_io_reader_descriptor
-#undef mlx_io_reader_tostring
-#undef mlx_io_reader_free
-#undef mlx_io_writer_new
-#undef mlx_io_writer_descriptor
-#undef mlx_io_writer_tostring
-#undef mlx_io_writer_free
-#undef mlx_linalg_cholesky
-#undef mlx_linalg_cholesky_inv
-#undef mlx_linalg_cross
-#undef mlx_linalg_eig
-#undef mlx_linalg_eigh
-#undef mlx_linalg_eigvals
-#undef mlx_linalg_eigvalsh
-#undef mlx_linalg_inv
-#undef mlx_linalg_lu
-#undef mlx_linalg_lu_factor
-#undef mlx_linalg_norm
-#undef mlx_linalg_norm_matrix
-#undef mlx_linalg_norm_l2
-#undef mlx_linalg_pinv
-#undef mlx_linalg_qr
-#undef mlx_linalg_solve
-#undef mlx_linalg_solve_triangular
-#undef mlx_linalg_svd
-#undef mlx_linalg_tri_inv
-#undef mlx_map_string_to_array_new
-#undef mlx_map_string_to_array_set
-#undef mlx_map_string_to_array_free
-#undef mlx_map_string_to_array_insert
-#undef mlx_map_string_to_array_get
-#undef mlx_map_string_to_array_iterator_new
-#undef mlx_map_string_to_array_iterator_free
-#undef mlx_map_string_to_array_iterator_next
-#undef mlx_map_string_to_string_new
-#undef mlx_map_string_to_string_set
-#undef mlx_map_string_to_string_free
-#undef mlx_map_string_to_string_insert
-#undef mlx_map_string_to_string_get
-#undef mlx_map_string_to_string_iterator_new
-#undef mlx_map_string_to_string_iterator_free
-#undef mlx_map_string_to_string_iterator_next
-#undef mlx_clear_cache
-#undef mlx_get_active_memory
-#undef mlx_get_cache_memory
-#undef mlx_get_memory_limit
-#undef mlx_get_peak_memory
-#undef mlx_reset_peak_memory
-#undef mlx_set_cache_limit
-#undef mlx_set_memory_limit
-#undef mlx_set_wired_limit
-#undef mlx_metal_device_info
-#undef mlx_metal_is_available
-#undef mlx_metal_start_capture
-#undef mlx_metal_stop_capture
-#undef mlx_abs
-#undef mlx_add
-#undef mlx_addmm
-#undef mlx_all_axes
-#undef mlx_all_axis
-#undef mlx_all
-#undef mlx_allclose
-#undef mlx_any_axes
-#undef mlx_any_axis
-#undef mlx_any
-#undef mlx_arange
-#undef mlx_arccos
-#undef mlx_arccosh
-#undef mlx_arcsin
-#undef mlx_arcsinh
-#undef mlx_arctan
-#undef mlx_arctan2
-#undef mlx_arctanh
-#undef mlx_argmax_axis
-#undef mlx_argmax
-#undef mlx_argmin_axis
-#undef mlx_argmin
-#undef mlx_argpartition_axis
-#undef mlx_argpartition
-#undef mlx_argsort_axis
-#undef mlx_argsort
-#undef mlx_array_equal
-#undef mlx_as_strided
-#undef mlx_astype
-#undef mlx_atleast_1d
-#undef mlx_atleast_2d
-#undef mlx_atleast_3d
-#undef mlx_bitwise_and
-#undef mlx_bitwise_invert
-#undef mlx_bitwise_or
-#undef mlx_bitwise_xor
-#undef mlx_block_masked_mm
-#undef mlx_broadcast_arrays
-#undef mlx_broadcast_to
-#undef mlx_ceil
-#undef mlx_clip
-#undef mlx_concatenate_axis
-#undef mlx_concatenate
-#undef mlx_conjugate
-#undef mlx_contiguous
-#undef mlx_conv1d
-#undef mlx_conv2d
-#undef mlx_conv3d
-#undef mlx_conv_general
-#undef mlx_conv_transpose1d
-#undef mlx_conv_transpose2d
-#undef mlx_conv_transpose3d
-#undef mlx_copy
-#undef mlx_cos
-#undef mlx_cosh
-#undef mlx_cummax
-#undef mlx_cummin
-#undef mlx_cumprod
-#undef mlx_cumsum
-#undef mlx_degrees
-#undef mlx_depends
-#undef mlx_dequantize
-#undef mlx_diag
-#undef mlx_diagonal
-#undef mlx_divide
-#undef mlx_divmod
-#undef mlx_einsum
-#undef mlx_equal
-#undef mlx_erf
-#undef mlx_erfinv
-#undef mlx_exp
-#undef mlx_expand_dims_axes
-#undef mlx_expand_dims
-#undef mlx_expm1
-#undef mlx_eye
-#undef mlx_flatten
-#undef mlx_floor
-#undef mlx_floor_divide
-#undef mlx_from_fp8
-#undef mlx_full
-#undef mlx_full_like
-#undef mlx_gather
-#undef mlx_gather_single
-#undef mlx_gather_mm
-#undef mlx_gather_qmm
-#undef mlx_greater
-#undef mlx_greater_equal
-#undef mlx_hadamard_transform
-#undef mlx_identity
-#undef mlx_imag
-#undef mlx_inner
-#undef mlx_isclose
-#undef mlx_isfinite
-#undef mlx_isinf
-#undef mlx_isnan
-#undef mlx_isneginf
-#undef mlx_isposinf
-#undef mlx_kron
-#undef mlx_left_shift
-#undef mlx_less
-#undef mlx_less_equal
-#undef mlx_linspace
-#undef mlx_log
-#undef mlx_log10
-#undef mlx_log1p
-#undef mlx_log2
-#undef mlx_logaddexp
-#undef mlx_logcumsumexp
-#undef mlx_logical_and
-#undef mlx_logical_not
-#undef mlx_logical_or
-#undef mlx_logsumexp_axes
-#undef mlx_logsumexp_axis
-#undef mlx_logsumexp
-#undef mlx_masked_scatter
-#undef mlx_matmul
-#undef mlx_max_axes
-#undef mlx_max_axis
-#undef mlx_max
-#undef mlx_maximum
-#undef mlx_mean_axes
-#undef mlx_mean_axis
-#undef mlx_mean
-#undef mlx_median
-#undef mlx_meshgrid
-#undef mlx_min_axes
-#undef mlx_min_axis
-#undef mlx_min
-#undef mlx_minimum
-#undef mlx_moveaxis
-#undef mlx_multiply
-#undef mlx_nan_to_num
-#undef mlx_negative
-#undef mlx_not_equal
-#undef mlx_number_of_elements
-#undef mlx_ones
-#undef mlx_ones_like
-#undef mlx_outer
-#undef mlx_pad
-#undef mlx_pad_symmetric
-#undef mlx_partition_axis
-#undef mlx_partition
-#undef mlx_power
-#undef mlx_prod_axes
-#undef mlx_prod_axis
-#undef mlx_prod
-#undef mlx_put_along_axis
-#undef mlx_qqmm
-#undef mlx_quantize
-#undef mlx_quantized_matmul
-#undef mlx_radians
-#undef mlx_real
-#undef mlx_reciprocal
-#undef mlx_remainder
-#undef mlx_repeat_axis
-#undef mlx_repeat
-#undef mlx_reshape
-#undef mlx_right_shift
-#undef mlx_roll_axis
-#undef mlx_roll_axes
-#undef mlx_roll
-#undef mlx_round
-#undef mlx_rsqrt
-#undef mlx_scatter
-#undef mlx_scatter_single
-#undef mlx_scatter_add
-#undef mlx_scatter_add_single
-#undef mlx_scatter_add_axis
-#undef mlx_scatter_max
-#undef mlx_scatter_max_single
-#undef mlx_scatter_min
-#undef mlx_scatter_min_single
-#undef mlx_scatter_prod
-#undef mlx_scatter_prod_single
-#undef mlx_segmented_mm
-#undef mlx_sigmoid
-#undef mlx_sign
-#undef mlx_sin
-#undef mlx_sinh
-#undef mlx_slice
-#undef mlx_slice_dynamic
-#undef mlx_slice_update
-#undef mlx_slice_update_dynamic
-#undef mlx_softmax_axes
-#undef mlx_softmax_axis
-#undef mlx_softmax
-#undef mlx_sort_axis
-#undef mlx_sort
-#undef mlx_split
-#undef mlx_split_sections
-#undef mlx_sqrt
-#undef mlx_square
-#undef mlx_squeeze_axes
-#undef mlx_squeeze_axis
-#undef mlx_squeeze
-#undef mlx_stack_axis
-#undef mlx_stack
-#undef mlx_std_axes
-#undef mlx_std_axis
-#undef mlx_std
-#undef mlx_stop_gradient
-#undef mlx_subtract
-#undef mlx_sum_axes
-#undef mlx_sum_axis
-#undef mlx_sum
-#undef mlx_swapaxes
-#undef mlx_take_axis
-#undef mlx_take
-#undef mlx_take_along_axis
-#undef mlx_tan
-#undef mlx_tanh
-#undef mlx_tensordot
-#undef mlx_tensordot_axis
-#undef mlx_tile
-#undef mlx_to_fp8
-#undef mlx_topk_axis
-#undef mlx_topk
-#undef mlx_trace
-#undef mlx_transpose_axes
-#undef mlx_transpose
-#undef mlx_tri
-#undef mlx_tril
-#undef mlx_triu
-#undef mlx_unflatten
-#undef mlx_var_axes
-#undef mlx_var_axis
-#undef mlx_var
-#undef mlx_view
-#undef mlx_where
-#undef mlx_zeros
-#undef mlx_zeros_like
-#undef mlx_random_bernoulli
-#undef mlx_random_bits
-#undef mlx_random_categorical_shape
-#undef mlx_random_categorical_num_samples
-#undef mlx_random_categorical
-#undef mlx_random_gumbel
-#undef mlx_random_key
-#undef mlx_random_laplace
-#undef mlx_random_multivariate_normal
-#undef mlx_random_normal_broadcast
-#undef mlx_random_normal
-#undef mlx_random_permutation
-#undef mlx_random_permutation_arange
-#undef mlx_random_randint
-#undef mlx_random_seed
-#undef mlx_random_split_num
-#undef mlx_random_split
-#undef mlx_random_truncated_normal
-#undef mlx_random_uniform
-#undef mlx_stream_new
-#undef mlx_stream_new_device
-#undef mlx_stream_set
-#undef mlx_stream_free
-#undef mlx_stream_tostring
-#undef mlx_stream_equal
-#undef mlx_stream_get_device
-#undef mlx_stream_get_index
-#undef mlx_synchronize
-#undef mlx_get_default_stream
-#undef mlx_set_default_stream
-#undef mlx_default_cpu_stream_new
-#undef mlx_default_gpu_stream_new
-#undef mlx_string_new
-#undef mlx_string_new_data
-#undef mlx_string_set
-#undef mlx_string_data
-#undef mlx_string_free
-#undef mlx_async_eval
-#undef mlx_checkpoint
-#undef mlx_custom_function
-#undef mlx_custom_vjp
-#undef mlx_eval
-#undef mlx_jvp
-#undef mlx_value_and_grad
-#undef mlx_vjp
-#undef mlx_detail_vmap_replace
-#undef mlx_detail_vmap_trace
-#undef mlx_vector_array_new
-#undef mlx_vector_array_set
-#undef mlx_vector_array_free
-#undef mlx_vector_array_new_data
-#undef mlx_vector_array_new_value
-#undef mlx_vector_array_set_data
-#undef mlx_vector_array_set_value
-#undef mlx_vector_array_append_data
-#undef mlx_vector_array_append_value
-#undef mlx_vector_array_size
-#undef mlx_vector_array_get
-#undef mlx_vector_vector_array_new
-#undef mlx_vector_vector_array_set
-#undef mlx_vector_vector_array_free
-#undef mlx_vector_vector_array_new_data
-#undef mlx_vector_vector_array_new_value
-#undef mlx_vector_vector_array_set_data
-#undef mlx_vector_vector_array_set_value
-#undef mlx_vector_vector_array_append_data
-#undef mlx_vector_vector_array_append_value
-#undef mlx_vector_vector_array_size
-#undef mlx_vector_vector_array_get
-#undef mlx_vector_int_new
-#undef mlx_vector_int_set
-#undef mlx_vector_int_free
-#undef mlx_vector_int_new_data
-#undef mlx_vector_int_new_value
-#undef mlx_vector_int_set_data
-#undef mlx_vector_int_set_value
-#undef mlx_vector_int_append_data
-#undef mlx_vector_int_append_value
-#undef mlx_vector_int_size
-#undef mlx_vector_int_get
-#undef mlx_vector_string_new
-#undef mlx_vector_string_set
-#undef mlx_vector_string_free
-#undef mlx_vector_string_new_data
-#undef mlx_vector_string_new_value
-#undef mlx_vector_string_set_data
-#undef mlx_vector_string_set_value
-#undef mlx_vector_string_append_data
-#undef mlx_vector_string_append_value
-#undef mlx_vector_string_size
-#undef mlx_vector_string_get
-#undef mlx_version
-
-// Function pointer declarations (defined in mlx.c, loaded via dlsym)
-extern size_t (*mlx_dtype_size_ptr)(mlx_dtype dtype);
-extern int (*mlx_array_tostring_ptr)(mlx_string* str, const mlx_array arr);
-extern mlx_array (*mlx_array_new_ptr)(void);
-extern int (*mlx_array_free_ptr)(mlx_array arr);
-extern mlx_array (*mlx_array_new_bool_ptr)(bool val);
-extern mlx_array (*mlx_array_new_int_ptr)(int val);
-extern mlx_array (*mlx_array_new_float32_ptr)(float val);
-extern mlx_array (*mlx_array_new_float_ptr)(float val);
-extern mlx_array (*mlx_array_new_float64_ptr)(double val);
-extern mlx_array (*mlx_array_new_double_ptr)(double val);
-extern mlx_array (*mlx_array_new_complex_ptr)(float real_val, float imag_val);
-extern mlx_array (*mlx_array_new_data_ptr)(const void* data, const int* shape, int dim, mlx_dtype dtype);
-extern int (*mlx_array_set_ptr)(mlx_array* arr, const mlx_array src);
-extern int (*mlx_array_set_bool_ptr)(mlx_array* arr, bool val);
-extern int (*mlx_array_set_int_ptr)(mlx_array* arr, int val);
-extern int (*mlx_array_set_float32_ptr)(mlx_array* arr, float val);
-extern int (*mlx_array_set_float_ptr)(mlx_array* arr, float val);
-extern int (*mlx_array_set_float64_ptr)(mlx_array* arr, double val);
-extern int (*mlx_array_set_double_ptr)(mlx_array* arr, double val);
-extern int (*mlx_array_set_complex_ptr)(mlx_array* arr, float real_val, float imag_val);
-extern int (*mlx_array_set_data_ptr)(mlx_array* arr, const void* data, const int* shape, int dim, mlx_dtype dtype);
-extern size_t (*mlx_array_itemsize_ptr)(const mlx_array arr);
-extern size_t (*mlx_array_size_ptr)(const mlx_array arr);
-extern size_t (*mlx_array_nbytes_ptr)(const mlx_array arr);
-extern size_t (*mlx_array_ndim_ptr)(const mlx_array arr);
-extern const int* (*mlx_array_shape_ptr)(const mlx_array arr);
-extern const size_t* (*mlx_array_strides_ptr)(const mlx_array arr);
-extern int (*mlx_array_dim_ptr)(const mlx_array arr, int dim);
-extern mlx_dtype (*mlx_array_dtype_ptr)(const mlx_array arr);
-extern int (*mlx_array_eval_ptr)(mlx_array arr);
-extern int (*mlx_array_item_bool_ptr)(bool* res, const mlx_array arr);
-extern int (*mlx_array_item_uint8_ptr)(uint8_t* res, const mlx_array arr);
-extern int (*mlx_array_item_uint16_ptr)(uint16_t* res, const mlx_array arr);
-extern int (*mlx_array_item_uint32_ptr)(uint32_t* res, const mlx_array arr);
-extern int (*mlx_array_item_uint64_ptr)(uint64_t* res, const mlx_array arr);
-extern int (*mlx_array_item_int8_ptr)(int8_t* res, const mlx_array arr);
-extern int (*mlx_array_item_int16_ptr)(int16_t* res, const mlx_array arr);
-extern int (*mlx_array_item_int32_ptr)(int32_t* res, const mlx_array arr);
-extern int (*mlx_array_item_int64_ptr)(int64_t* res, const mlx_array arr);
-extern int (*mlx_array_item_float32_ptr)(float* res, const mlx_array arr);
-extern int (*mlx_array_item_float64_ptr)(double* res, const mlx_array arr);
-extern int (*mlx_array_item_complex64_ptr)(float _Complex* res, const mlx_array arr);
-#if defined(__aarch64__) || defined(_M_ARM64)
-extern int (*mlx_array_item_float16_ptr)(float16_t* res, const mlx_array arr);
-#endif
-#if defined(__aarch64__) || defined(_M_ARM64)
-extern int (*mlx_array_item_bfloat16_ptr)(bfloat16_t* res, const mlx_array arr);
-#endif
-extern const bool* (*mlx_array_data_bool_ptr)(const mlx_array arr);
-extern const uint8_t* (*mlx_array_data_uint8_ptr)(const mlx_array arr);
-extern const uint16_t* (*mlx_array_data_uint16_ptr)(const mlx_array arr);
-extern const uint32_t* (*mlx_array_data_uint32_ptr)(const mlx_array arr);
-extern const uint64_t* (*mlx_array_data_uint64_ptr)(const mlx_array arr);
-extern const int8_t* (*mlx_array_data_int8_ptr)(const mlx_array arr);
-extern const int16_t* (*mlx_array_data_int16_ptr)(const mlx_array arr);
-extern const int32_t* (*mlx_array_data_int32_ptr)(const mlx_array arr);
-extern const int64_t* (*mlx_array_data_int64_ptr)(const mlx_array arr);
-extern const float* (*mlx_array_data_float32_ptr)(const mlx_array arr);
-extern const double* (*mlx_array_data_float64_ptr)(const mlx_array arr);
-extern const float _Complex* (*mlx_array_data_complex64_ptr)(const mlx_array arr);
-#if defined(__aarch64__) || defined(_M_ARM64)
-extern const float16_t* (*mlx_array_data_float16_ptr)(const mlx_array arr);
-#endif
-#if defined(__aarch64__) || defined(_M_ARM64)
-extern const bfloat16_t* (*mlx_array_data_bfloat16_ptr)(const mlx_array arr);
-#endif
-extern int (*_mlx_array_is_available_ptr)(bool* res, const mlx_array arr);
-extern int (*_mlx_array_wait_ptr)(const mlx_array arr);
-extern int (*_mlx_array_is_contiguous_ptr)(bool* res, const mlx_array arr);
-extern int (*_mlx_array_is_row_contiguous_ptr)(bool* res, const mlx_array arr);
-extern int (*_mlx_array_is_col_contiguous_ptr)(bool* res, const mlx_array arr);
-extern mlx_closure (*mlx_closure_new_ptr)(void);
-extern int (*mlx_closure_free_ptr)(mlx_closure cls);
-extern mlx_closure (*mlx_closure_new_func_ptr)(int (*fun)(mlx_vector_array*, const mlx_vector_array));
-extern mlx_closure (*mlx_closure_new_func_payload_ptr)(int (*fun)(mlx_vector_array*, const mlx_vector_array, void*), void* payload, void (*dtor)(void*));
-extern int (*mlx_closure_set_ptr)(mlx_closure* cls, const mlx_closure src);
-extern int (*mlx_closure_apply_ptr)(mlx_vector_array* res, mlx_closure cls, const mlx_vector_array input);
-extern mlx_closure (*mlx_closure_new_unary_ptr)(int (*fun)(mlx_array*, const mlx_array));
-extern mlx_closure_kwargs (*mlx_closure_kwargs_new_ptr)(void);
-extern int (*mlx_closure_kwargs_free_ptr)(mlx_closure_kwargs cls);
-extern mlx_closure_kwargs (*mlx_closure_kwargs_new_func_ptr)(int (*fun)( mlx_vector_array*, const mlx_vector_array, const mlx_map_string_to_array));
-extern mlx_closure_kwargs (*mlx_closure_kwargs_new_func_payload_ptr)(int (*fun)( mlx_vector_array*, const mlx_vector_array, const mlx_map_string_to_array, void*), void* payload, void (*dtor)(void*));
-extern int (*mlx_closure_kwargs_set_ptr)(mlx_closure_kwargs* cls, const mlx_closure_kwargs src);
-extern int (*mlx_closure_kwargs_apply_ptr)(mlx_vector_array* res, mlx_closure_kwargs cls, const mlx_vector_array input_0, const mlx_map_string_to_array input_1);
-extern mlx_closure_value_and_grad (*mlx_closure_value_and_grad_new_ptr)(void);
-extern int (*mlx_closure_value_and_grad_free_ptr)(mlx_closure_value_and_grad cls);
-extern mlx_closure_value_and_grad (*mlx_closure_value_and_grad_new_func_ptr)(int (*fun)(mlx_vector_array*, mlx_vector_array*, const mlx_vector_array));
-extern mlx_closure_value_and_grad (*mlx_closure_value_and_grad_new_func_payload_ptr)(int (*fun)( mlx_vector_array*, mlx_vector_array*, const mlx_vector_array, void*), void* payload, void (*dtor)(void*));
-extern int (*mlx_closure_value_and_grad_set_ptr)(mlx_closure_value_and_grad* cls, const mlx_closure_value_and_grad src);
-extern int (*mlx_closure_value_and_grad_apply_ptr)(mlx_vector_array* res_0, mlx_vector_array* res_1, mlx_closure_value_and_grad cls, const mlx_vector_array input);
-extern mlx_closure_custom (*mlx_closure_custom_new_ptr)(void);
-extern int (*mlx_closure_custom_free_ptr)(mlx_closure_custom cls);
-extern mlx_closure_custom (*mlx_closure_custom_new_func_ptr)(int (*fun)( mlx_vector_array*, const mlx_vector_array, const mlx_vector_array, const mlx_vector_array));
-extern mlx_closure_custom (*mlx_closure_custom_new_func_payload_ptr)(int (*fun)( mlx_vector_array*, const mlx_vector_array, const mlx_vector_array, const mlx_vector_array, void*), void* payload, void (*dtor)(void*));
-extern int (*mlx_closure_custom_set_ptr)(mlx_closure_custom* cls, const mlx_closure_custom src);
-extern int (*mlx_closure_custom_apply_ptr)(mlx_vector_array* res, mlx_closure_custom cls, const mlx_vector_array input_0, const mlx_vector_array input_1, const mlx_vector_array input_2);
-extern mlx_closure_custom_jvp (*mlx_closure_custom_jvp_new_ptr)(void);
-extern int (*mlx_closure_custom_jvp_free_ptr)(mlx_closure_custom_jvp cls);
-extern mlx_closure_custom_jvp (*mlx_closure_custom_jvp_new_func_ptr)(int (*fun)( mlx_vector_array*, const mlx_vector_array, const mlx_vector_array, const int*, size_t _num));
-extern mlx_closure_custom_jvp (*mlx_closure_custom_jvp_new_func_payload_ptr)(int (*fun)( mlx_vector_array*, const mlx_vector_array, const mlx_vector_array, const int*, size_t _num, void*), void* payload, void (*dtor)(void*));
-extern int (*mlx_closure_custom_jvp_set_ptr)(mlx_closure_custom_jvp* cls, const mlx_closure_custom_jvp src);
-extern int (*mlx_closure_custom_jvp_apply_ptr)(mlx_vector_array* res, mlx_closure_custom_jvp cls, const mlx_vector_array input_0, const mlx_vector_array input_1, const int* input_2, size_t input_2_num);
-extern mlx_closure_custom_vmap (*mlx_closure_custom_vmap_new_ptr)(void);
-extern int (*mlx_closure_custom_vmap_free_ptr)(mlx_closure_custom_vmap cls);
-extern mlx_closure_custom_vmap (*mlx_closure_custom_vmap_new_func_ptr)(int (*fun)( mlx_vector_array*, mlx_vector_int*, const mlx_vector_array, const int*, size_t _num));
-extern mlx_closure_custom_vmap (*mlx_closure_custom_vmap_new_func_payload_ptr)(int (*fun)( mlx_vector_array*, mlx_vector_int*, const mlx_vector_array, const int*, size_t _num, void*), void* payload, void (*dtor)(void*));
-extern int (*mlx_closure_custom_vmap_set_ptr)(mlx_closure_custom_vmap* cls, const mlx_closure_custom_vmap src);
-extern int (*mlx_closure_custom_vmap_apply_ptr)(mlx_vector_array* res_0, mlx_vector_int* res_1, mlx_closure_custom_vmap cls, const mlx_vector_array input_0, const int* input_1, size_t input_1_num);
-extern int (*mlx_compile_ptr)(mlx_closure* res, const mlx_closure fun, bool shapeless);
-extern int (*mlx_detail_compile_ptr)(mlx_closure* res, const mlx_closure fun, uintptr_t fun_id, bool shapeless, const uint64_t* constants, size_t constants_num);
-extern int (*mlx_detail_compile_clear_cache_ptr)(void);
-extern int (*mlx_detail_compile_erase_ptr)(uintptr_t fun_id);
-extern int (*mlx_disable_compile_ptr)(void);
-extern int (*mlx_enable_compile_ptr)(void);
-extern int (*mlx_set_compile_mode_ptr)(mlx_compile_mode mode);
-extern mlx_device (*mlx_device_new_ptr)(void);
-extern mlx_device (*mlx_device_new_type_ptr)(mlx_device_type type, int index);
-extern int (*mlx_device_free_ptr)(mlx_device dev);
-extern int (*mlx_device_set_ptr)(mlx_device* dev, const mlx_device src);
-extern int (*mlx_device_tostring_ptr)(mlx_string* str, mlx_device dev);
-extern bool (*mlx_device_equal_ptr)(mlx_device lhs, mlx_device rhs);
-extern int (*mlx_device_get_index_ptr)(int* index, mlx_device dev);
-extern int (*mlx_device_get_type_ptr)(mlx_device_type* type, mlx_device dev);
-extern int (*mlx_get_default_device_ptr)(mlx_device* dev);
-extern int (*mlx_set_default_device_ptr)(mlx_device dev);
-extern int (*mlx_distributed_all_gather_ptr)(mlx_array* res, const mlx_array x, const mlx_distributed_group group , const mlx_stream S);
-extern int (*mlx_distributed_all_max_ptr)(mlx_array* res, const mlx_array x, const mlx_distributed_group group , const mlx_stream s);
-extern int (*mlx_distributed_all_min_ptr)(mlx_array* res, const mlx_array x, const mlx_distributed_group group , const mlx_stream s);
-extern int (*mlx_distributed_all_sum_ptr)(mlx_array* res, const mlx_array x, const mlx_distributed_group group , const mlx_stream s);
-extern int (*mlx_distributed_recv_ptr)(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, int src, const mlx_distributed_group group , const mlx_stream s);
-extern int (*mlx_distributed_recv_like_ptr)(mlx_array* res, const mlx_array x, int src, const mlx_distributed_group group , const mlx_stream s);
-extern int (*mlx_distributed_send_ptr)(mlx_array* res, const mlx_array x, int dst, const mlx_distributed_group group , const mlx_stream s);
-extern int (*mlx_distributed_sum_scatter_ptr)(mlx_array* res, const mlx_array x, const mlx_distributed_group group , const mlx_stream s);
-extern int (*mlx_distributed_group_rank_ptr)(mlx_distributed_group group);
-extern int (*mlx_distributed_group_size_ptr)(mlx_distributed_group group);
-extern mlx_distributed_group (*mlx_distributed_group_split_ptr)(mlx_distributed_group group, int color, int key);
-extern bool (*mlx_distributed_is_available_ptr)(void);
-extern mlx_distributed_group (*mlx_distributed_init_ptr)(bool strict);
-extern void (*mlx_set_error_handler_ptr)(mlx_error_handler_func handler, void* data, void (*dtor)(void*));
-extern void (*_mlx_error_ptr)(const char* file, const int line, const char* fmt, ...);
-extern int (*mlx_export_function_ptr)(const char* file, const mlx_closure fun, const mlx_vector_array args, bool shapeless);
-extern int (*mlx_export_function_kwargs_ptr)(const char* file, const mlx_closure_kwargs fun, const mlx_vector_array args, const mlx_map_string_to_array kwargs, bool shapeless);
-extern mlx_function_exporter (*mlx_function_exporter_new_ptr)(const char* file, const mlx_closure fun, bool shapeless);
-extern int (*mlx_function_exporter_free_ptr)(mlx_function_exporter xfunc);
-extern int (*mlx_function_exporter_apply_ptr)(const mlx_function_exporter xfunc, const mlx_vector_array args);
-extern int (*mlx_function_exporter_apply_kwargs_ptr)(const mlx_function_exporter xfunc, const mlx_vector_array args, const mlx_map_string_to_array kwargs);
-extern mlx_imported_function (*mlx_imported_function_new_ptr)(const char* file);
-extern int (*mlx_imported_function_free_ptr)(mlx_imported_function xfunc);
-extern int (*mlx_imported_function_apply_ptr)(mlx_vector_array* res, const mlx_imported_function xfunc, const mlx_vector_array args);
-extern int (*mlx_imported_function_apply_kwargs_ptr)(mlx_vector_array* res, const mlx_imported_function xfunc, const mlx_vector_array args, const mlx_map_string_to_array kwargs);
-extern mlx_fast_cuda_kernel_config (*mlx_fast_cuda_kernel_config_new_ptr)(void);
-extern void (*mlx_fast_cuda_kernel_config_free_ptr)(mlx_fast_cuda_kernel_config cls);
-extern int (*mlx_fast_cuda_kernel_config_add_output_arg_ptr)(mlx_fast_cuda_kernel_config cls, const int* shape, size_t size, mlx_dtype dtype);
-extern int (*mlx_fast_cuda_kernel_config_set_grid_ptr)(mlx_fast_cuda_kernel_config cls, int grid1, int grid2, int grid3);
-extern int (*mlx_fast_cuda_kernel_config_set_thread_group_ptr)(mlx_fast_cuda_kernel_config cls, int thread1, int thread2, int thread3);
-extern int (*mlx_fast_cuda_kernel_config_set_init_value_ptr)(mlx_fast_cuda_kernel_config cls, float value);
-extern int (*mlx_fast_cuda_kernel_config_set_verbose_ptr)(mlx_fast_cuda_kernel_config cls, bool verbose);
-extern int (*mlx_fast_cuda_kernel_config_add_template_arg_dtype_ptr)(mlx_fast_cuda_kernel_config cls, const char* name, mlx_dtype dtype);
-extern int (*mlx_fast_cuda_kernel_config_add_template_arg_int_ptr)(mlx_fast_cuda_kernel_config cls, const char* name, int value);
-extern int (*mlx_fast_cuda_kernel_config_add_template_arg_bool_ptr)(mlx_fast_cuda_kernel_config cls, const char* name, bool value);
-extern mlx_fast_cuda_kernel (*mlx_fast_cuda_kernel_new_ptr)(const char* name, const mlx_vector_string input_names, const mlx_vector_string output_names, const char* source, const char* header, bool ensure_row_contiguous, int shared_memory);
-extern void (*mlx_fast_cuda_kernel_free_ptr)(mlx_fast_cuda_kernel cls);
-extern int (*mlx_fast_cuda_kernel_apply_ptr)(mlx_vector_array* outputs, mlx_fast_cuda_kernel cls, const mlx_vector_array inputs, const mlx_fast_cuda_kernel_config config, const mlx_stream stream);
-extern int (*mlx_fast_layer_norm_ptr)(mlx_array* res, const mlx_array x, const mlx_array weight , const mlx_array bias , float eps, const mlx_stream s);
-extern mlx_fast_metal_kernel_config (*mlx_fast_metal_kernel_config_new_ptr)(void);
-extern void (*mlx_fast_metal_kernel_config_free_ptr)(mlx_fast_metal_kernel_config cls);
-extern int (*mlx_fast_metal_kernel_config_add_output_arg_ptr)(mlx_fast_metal_kernel_config cls, const int* shape, size_t size, mlx_dtype dtype);
-extern int (*mlx_fast_metal_kernel_config_set_grid_ptr)(mlx_fast_metal_kernel_config cls, int grid1, int grid2, int grid3);
-extern int (*mlx_fast_metal_kernel_config_set_thread_group_ptr)(mlx_fast_metal_kernel_config cls, int thread1, int thread2, int thread3);
-extern int (*mlx_fast_metal_kernel_config_set_init_value_ptr)(mlx_fast_metal_kernel_config cls, float value);
-extern int (*mlx_fast_metal_kernel_config_set_verbose_ptr)(mlx_fast_metal_kernel_config cls, bool verbose);
-extern int (*mlx_fast_metal_kernel_config_add_template_arg_dtype_ptr)(mlx_fast_metal_kernel_config cls, const char* name, mlx_dtype dtype);
-extern int (*mlx_fast_metal_kernel_config_add_template_arg_int_ptr)(mlx_fast_metal_kernel_config cls, const char* name, int value);
-extern int (*mlx_fast_metal_kernel_config_add_template_arg_bool_ptr)(mlx_fast_metal_kernel_config cls, const char* name, bool value);
-extern mlx_fast_metal_kernel (*mlx_fast_metal_kernel_new_ptr)(const char* name, const mlx_vector_string input_names, const mlx_vector_string output_names, const char* source, const char* header, bool ensure_row_contiguous, bool atomic_outputs);
-extern void (*mlx_fast_metal_kernel_free_ptr)(mlx_fast_metal_kernel cls);
-extern int (*mlx_fast_metal_kernel_apply_ptr)(mlx_vector_array* outputs, mlx_fast_metal_kernel cls, const mlx_vector_array inputs, const mlx_fast_metal_kernel_config config, const mlx_stream stream);
-extern int (*mlx_fast_rms_norm_ptr)(mlx_array* res, const mlx_array x, const mlx_array weight , float eps, const mlx_stream s);
-extern int (*mlx_fast_rope_ptr)(mlx_array* res, const mlx_array x, int dims, bool traditional, mlx_optional_float base, float scale, int offset, const mlx_array freqs , const mlx_stream s);
-extern int (*mlx_fast_rope_dynamic_ptr)(mlx_array* res, const mlx_array x, int dims, bool traditional, mlx_optional_float base, float scale, const mlx_array offset, const mlx_array freqs , const mlx_stream s);
-extern int (*mlx_fast_scaled_dot_product_attention_ptr)(mlx_array* res, const mlx_array queries, const mlx_array keys, const mlx_array values, float scale, const char* mask_mode, const mlx_array mask_arr , const mlx_array sinks , const mlx_stream s);
-extern int (*mlx_fft_fft_ptr)(mlx_array* res, const mlx_array a, int n, int axis, const mlx_stream s);
-extern int (*mlx_fft_fft2_ptr)(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s);
-extern int (*mlx_fft_fftn_ptr)(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s);
-extern int (*mlx_fft_fftshift_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, const mlx_stream s);
-extern int (*mlx_fft_ifft_ptr)(mlx_array* res, const mlx_array a, int n, int axis, const mlx_stream s);
-extern int (*mlx_fft_ifft2_ptr)(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s);
-extern int (*mlx_fft_ifftn_ptr)(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s);
-extern int (*mlx_fft_ifftshift_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, const mlx_stream s);
-extern int (*mlx_fft_irfft_ptr)(mlx_array* res, const mlx_array a, int n, int axis, const mlx_stream s);
-extern int (*mlx_fft_irfft2_ptr)(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s);
-extern int (*mlx_fft_irfftn_ptr)(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s);
-extern int (*mlx_fft_rfft_ptr)(mlx_array* res, const mlx_array a, int n, int axis, const mlx_stream s);
-extern int (*mlx_fft_rfft2_ptr)(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s);
-extern int (*mlx_fft_rfftn_ptr)(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s);
-extern int (*mlx_load_reader_ptr)(mlx_array* res, mlx_io_reader in_stream, const mlx_stream s);
-extern int (*mlx_load_ptr)(mlx_array* res, const char* file, const mlx_stream s);
-extern int (*mlx_load_safetensors_reader_ptr)(mlx_map_string_to_array* res_0, mlx_map_string_to_string* res_1, mlx_io_reader in_stream, const mlx_stream s);
-extern int (*mlx_load_safetensors_ptr)(mlx_map_string_to_array* res_0, mlx_map_string_to_string* res_1, const char* file, const mlx_stream s);
-extern int (*mlx_save_writer_ptr)(mlx_io_writer out_stream, const mlx_array a);
-extern int (*mlx_save_ptr)(const char* file, const mlx_array a);
-extern int (*mlx_save_safetensors_writer_ptr)(mlx_io_writer in_stream, const mlx_map_string_to_array param, const mlx_map_string_to_string metadata);
-extern int (*mlx_save_safetensors_ptr)(const char* file, const mlx_map_string_to_array param, const mlx_map_string_to_string metadata);
-extern mlx_io_reader (*mlx_io_reader_new_ptr)(void* desc, mlx_io_vtable vtable);
-extern int (*mlx_io_reader_descriptor_ptr)(void** desc_, mlx_io_reader io);
-extern int (*mlx_io_reader_tostring_ptr)(mlx_string* str_, mlx_io_reader io);
-extern int (*mlx_io_reader_free_ptr)(mlx_io_reader io);
-extern mlx_io_writer (*mlx_io_writer_new_ptr)(void* desc, mlx_io_vtable vtable);
-extern int (*mlx_io_writer_descriptor_ptr)(void** desc_, mlx_io_writer io);
-extern int (*mlx_io_writer_tostring_ptr)(mlx_string* str_, mlx_io_writer io);
-extern int (*mlx_io_writer_free_ptr)(mlx_io_writer io);
-extern int (*mlx_linalg_cholesky_ptr)(mlx_array* res, const mlx_array a, bool upper, const mlx_stream s);
-extern int (*mlx_linalg_cholesky_inv_ptr)(mlx_array* res, const mlx_array a, bool upper, const mlx_stream s);
-extern int (*mlx_linalg_cross_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, int axis, const mlx_stream s);
-extern int (*mlx_linalg_eig_ptr)(mlx_array* res_0, mlx_array* res_1, const mlx_array a, const mlx_stream s);
-extern int (*mlx_linalg_eigh_ptr)(mlx_array* res_0, mlx_array* res_1, const mlx_array a, const char* UPLO, const mlx_stream s);
-extern int (*mlx_linalg_eigvals_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_linalg_eigvalsh_ptr)(mlx_array* res, const mlx_array a, const char* UPLO, const mlx_stream s);
-extern int (*mlx_linalg_inv_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_linalg_lu_ptr)(mlx_vector_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_linalg_lu_factor_ptr)(mlx_array* res_0, mlx_array* res_1, const mlx_array a, const mlx_stream s);
-extern int (*mlx_linalg_norm_ptr)(mlx_array* res, const mlx_array a, double ord, const int* axis , size_t axis_num, bool keepdims, const mlx_stream s);
-extern int (*mlx_linalg_norm_matrix_ptr)(mlx_array* res, const mlx_array a, const char* ord, const int* axis , size_t axis_num, bool keepdims, const mlx_stream s);
-extern int (*mlx_linalg_norm_l2_ptr)(mlx_array* res, const mlx_array a, const int* axis , size_t axis_num, bool keepdims, const mlx_stream s);
-extern int (*mlx_linalg_pinv_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_linalg_qr_ptr)(mlx_array* res_0, mlx_array* res_1, const mlx_array a, const mlx_stream s);
-extern int (*mlx_linalg_solve_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_linalg_solve_triangular_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, bool upper, const mlx_stream s);
-extern int (*mlx_linalg_svd_ptr)(mlx_vector_array* res, const mlx_array a, bool compute_uv, const mlx_stream s);
-extern int (*mlx_linalg_tri_inv_ptr)(mlx_array* res, const mlx_array a, bool upper, const mlx_stream s);
-extern mlx_map_string_to_array (*mlx_map_string_to_array_new_ptr)(void);
-extern int (*mlx_map_string_to_array_set_ptr)(mlx_map_string_to_array* map, const mlx_map_string_to_array src);
-extern int (*mlx_map_string_to_array_free_ptr)(mlx_map_string_to_array map);
-extern int (*mlx_map_string_to_array_insert_ptr)(mlx_map_string_to_array map, const char* key, const mlx_array value);
-extern int (*mlx_map_string_to_array_get_ptr)(mlx_array* value, const mlx_map_string_to_array map, const char* key);
-extern mlx_map_string_to_array_iterator (*mlx_map_string_to_array_iterator_new_ptr)(mlx_map_string_to_array map);
-extern int (*mlx_map_string_to_array_iterator_free_ptr)(mlx_map_string_to_array_iterator it);
-extern int (*mlx_map_string_to_array_iterator_next_ptr)(const char** key, mlx_array* value, mlx_map_string_to_array_iterator it);
-extern mlx_map_string_to_string (*mlx_map_string_to_string_new_ptr)(void);
-extern int (*mlx_map_string_to_string_set_ptr)(mlx_map_string_to_string* map, const mlx_map_string_to_string src);
-extern int (*mlx_map_string_to_string_free_ptr)(mlx_map_string_to_string map);
-extern int (*mlx_map_string_to_string_insert_ptr)(mlx_map_string_to_string map, const char* key, const char* value);
-extern int (*mlx_map_string_to_string_get_ptr)(const char** value, const mlx_map_string_to_string map, const char* key);
-extern mlx_map_string_to_string_iterator (*mlx_map_string_to_string_iterator_new_ptr)(mlx_map_string_to_string map);
-extern int (*mlx_map_string_to_string_iterator_free_ptr)(mlx_map_string_to_string_iterator it);
-extern int (*mlx_map_string_to_string_iterator_next_ptr)(const char** key, const char** value, mlx_map_string_to_string_iterator it);
-extern int (*mlx_clear_cache_ptr)(void);
-extern int (*mlx_get_active_memory_ptr)(size_t* res);
-extern int (*mlx_get_cache_memory_ptr)(size_t* res);
-extern int (*mlx_get_memory_limit_ptr)(size_t* res);
-extern int (*mlx_get_peak_memory_ptr)(size_t* res);
-extern int (*mlx_reset_peak_memory_ptr)(void);
-extern int (*mlx_set_cache_limit_ptr)(size_t* res, size_t limit);
-extern int (*mlx_set_memory_limit_ptr)(size_t* res, size_t limit);
-extern int (*mlx_set_wired_limit_ptr)(size_t* res, size_t limit);
-extern mlx_metal_device_info_t (*mlx_metal_device_info_ptr)(void);
-extern int (*mlx_metal_is_available_ptr)(bool* res);
-extern int (*mlx_metal_start_capture_ptr)(const char* path);
-extern int (*mlx_metal_stop_capture_ptr)(void);
-extern int (*mlx_abs_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_add_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_addmm_ptr)(mlx_array* res, const mlx_array c, const mlx_array a, const mlx_array b, float alpha, float beta, const mlx_stream s);
-extern int (*mlx_all_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s);
-extern int (*mlx_all_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s);
-extern int (*mlx_all_ptr)(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s);
-extern int (*mlx_allclose_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, double rtol, double atol, bool equal_nan, const mlx_stream s);
-extern int (*mlx_any_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s);
-extern int (*mlx_any_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s);
-extern int (*mlx_any_ptr)(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s);
-extern int (*mlx_arange_ptr)(mlx_array* res, double start, double stop, double step, mlx_dtype dtype, const mlx_stream s);
-extern int (*mlx_arccos_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_arccosh_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_arcsin_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_arcsinh_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_arctan_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_arctan2_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_arctanh_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_argmax_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s);
-extern int (*mlx_argmax_ptr)(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s);
-extern int (*mlx_argmin_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s);
-extern int (*mlx_argmin_ptr)(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s);
-extern int (*mlx_argpartition_axis_ptr)(mlx_array* res, const mlx_array a, int kth, int axis, const mlx_stream s);
-extern int (*mlx_argpartition_ptr)(mlx_array* res, const mlx_array a, int kth, const mlx_stream s);
-extern int (*mlx_argsort_axis_ptr)(mlx_array* res, const mlx_array a, int axis, const mlx_stream s);
-extern int (*mlx_argsort_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_array_equal_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, bool equal_nan, const mlx_stream s);
-extern int (*mlx_as_strided_ptr)(mlx_array* res, const mlx_array a, const int* shape, size_t shape_num, const int64_t* strides, size_t strides_num, size_t offset, const mlx_stream s);
-extern int (*mlx_astype_ptr)(mlx_array* res, const mlx_array a, mlx_dtype dtype, const mlx_stream s);
-extern int (*mlx_atleast_1d_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_atleast_2d_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_atleast_3d_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_bitwise_and_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_bitwise_invert_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_bitwise_or_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_bitwise_xor_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_block_masked_mm_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, int block_size, const mlx_array mask_out , const mlx_array mask_lhs , const mlx_array mask_rhs , const mlx_stream s);
-extern int (*mlx_broadcast_arrays_ptr)(mlx_vector_array* res, const mlx_vector_array inputs, const mlx_stream s);
-extern int (*mlx_broadcast_to_ptr)(mlx_array* res, const mlx_array a, const int* shape, size_t shape_num, const mlx_stream s);
-extern int (*mlx_ceil_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_clip_ptr)(mlx_array* res, const mlx_array a, const mlx_array a_min , const mlx_array a_max , const mlx_stream s);
-extern int (*mlx_concatenate_axis_ptr)(mlx_array* res, const mlx_vector_array arrays, int axis, const mlx_stream s);
-extern int (*mlx_concatenate_ptr)(mlx_array* res, const mlx_vector_array arrays, const mlx_stream s);
-extern int (*mlx_conjugate_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_contiguous_ptr)(mlx_array* res, const mlx_array a, bool allow_col_major, const mlx_stream s);
-extern int (*mlx_conv1d_ptr)(mlx_array* res, const mlx_array input, const mlx_array weight, int stride, int padding, int dilation, int groups, const mlx_stream s);
-extern int (*mlx_conv2d_ptr)(mlx_array* res, const mlx_array input, const mlx_array weight, int stride_0, int stride_1, int padding_0, int padding_1, int dilation_0, int dilation_1, int groups, const mlx_stream s);
-extern int (*mlx_conv3d_ptr)(mlx_array* res, const mlx_array input, const mlx_array weight, int stride_0, int stride_1, int stride_2, int padding_0, int padding_1, int padding_2, int dilation_0, int dilation_1, int dilation_2, int groups, const mlx_stream s);
-extern int (*mlx_conv_general_ptr)(mlx_array* res, const mlx_array input, const mlx_array weight, const int* stride, size_t stride_num, const int* padding_lo, size_t padding_lo_num, const int* padding_hi, size_t padding_hi_num, const int* kernel_dilation, size_t kernel_dilation_num, const int* input_dilation, size_t input_dilation_num, int groups, bool flip, const mlx_stream s);
-extern int (*mlx_conv_transpose1d_ptr)(mlx_array* res, const mlx_array input, const mlx_array weight, int stride, int padding, int dilation, int output_padding, int groups, const mlx_stream s);
-extern int (*mlx_conv_transpose2d_ptr)(mlx_array* res, const mlx_array input, const mlx_array weight, int stride_0, int stride_1, int padding_0, int padding_1, int dilation_0, int dilation_1, int output_padding_0, int output_padding_1, int groups, const mlx_stream s);
-extern int (*mlx_conv_transpose3d_ptr)(mlx_array* res, const mlx_array input, const mlx_array weight, int stride_0, int stride_1, int stride_2, int padding_0, int padding_1, int padding_2, int dilation_0, int dilation_1, int dilation_2, int output_padding_0, int output_padding_1, int output_padding_2, int groups, const mlx_stream s);
-extern int (*mlx_copy_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_cos_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_cosh_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_cummax_ptr)(mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s);
-extern int (*mlx_cummin_ptr)(mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s);
-extern int (*mlx_cumprod_ptr)(mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s);
-extern int (*mlx_cumsum_ptr)(mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s);
-extern int (*mlx_degrees_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_depends_ptr)(mlx_vector_array* res, const mlx_vector_array inputs, const mlx_vector_array dependencies);
-extern int (*mlx_dequantize_ptr)(mlx_array* res, const mlx_array w, const mlx_array scales, const mlx_array biases , mlx_optional_int group_size, mlx_optional_int bits, const char* mode, mlx_optional_dtype dtype, const mlx_stream s);
-extern int (*mlx_diag_ptr)(mlx_array* res, const mlx_array a, int k, const mlx_stream s);
-extern int (*mlx_diagonal_ptr)(mlx_array* res, const mlx_array a, int offset, int axis1, int axis2, const mlx_stream s);
-extern int (*mlx_divide_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_divmod_ptr)(mlx_vector_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_einsum_ptr)(mlx_array* res, const char* subscripts, const mlx_vector_array operands, const mlx_stream s);
-extern int (*mlx_equal_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_erf_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_erfinv_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_exp_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_expand_dims_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, const mlx_stream s);
-extern int (*mlx_expand_dims_ptr)(mlx_array* res, const mlx_array a, int axis, const mlx_stream s);
-extern int (*mlx_expm1_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_eye_ptr)(mlx_array* res, int n, int m, int k, mlx_dtype dtype, const mlx_stream s);
-extern int (*mlx_flatten_ptr)(mlx_array* res, const mlx_array a, int start_axis, int end_axis, const mlx_stream s);
-extern int (*mlx_floor_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_floor_divide_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_from_fp8_ptr)(mlx_array* res, const mlx_array x, mlx_dtype dtype, const mlx_stream s);
-extern int (*mlx_full_ptr)(mlx_array* res, const int* shape, size_t shape_num, const mlx_array vals, mlx_dtype dtype, const mlx_stream s);
-extern int (*mlx_full_like_ptr)(mlx_array* res, const mlx_array a, const mlx_array vals, mlx_dtype dtype, const mlx_stream s);
-extern int (*mlx_gather_ptr)(mlx_array* res, const mlx_array a, const mlx_vector_array indices, const int* axes, size_t axes_num, const int* slice_sizes, size_t slice_sizes_num, const mlx_stream s);
-extern int (*mlx_gather_single_ptr)(mlx_array* res, const mlx_array a, const mlx_array indices, int axis, const int* slice_sizes, size_t slice_sizes_num, const mlx_stream s);
-extern int (*mlx_gather_mm_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_array lhs_indices , const mlx_array rhs_indices , bool sorted_indices, const mlx_stream s);
-extern int (*mlx_gather_qmm_ptr)(mlx_array* res, const mlx_array x, const mlx_array w, const mlx_array scales, const mlx_array biases , const mlx_array lhs_indices , const mlx_array rhs_indices , bool transpose, mlx_optional_int group_size, mlx_optional_int bits, const char* mode, bool sorted_indices, const mlx_stream s);
-extern int (*mlx_greater_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_greater_equal_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_hadamard_transform_ptr)(mlx_array* res, const mlx_array a, mlx_optional_float scale, const mlx_stream s);
-extern int (*mlx_identity_ptr)(mlx_array* res, int n, mlx_dtype dtype, const mlx_stream s);
-extern int (*mlx_imag_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_inner_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_isclose_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, double rtol, double atol, bool equal_nan, const mlx_stream s);
-extern int (*mlx_isfinite_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_isinf_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_isnan_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_isneginf_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_isposinf_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_kron_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_left_shift_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_less_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_less_equal_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_linspace_ptr)(mlx_array* res, double start, double stop, int num, mlx_dtype dtype, const mlx_stream s);
-extern int (*mlx_log_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_log10_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_log1p_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_log2_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_logaddexp_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_logcumsumexp_ptr)(mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s);
-extern int (*mlx_logical_and_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_logical_not_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_logical_or_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_logsumexp_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s);
-extern int (*mlx_logsumexp_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s);
-extern int (*mlx_logsumexp_ptr)(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s);
-extern int (*mlx_masked_scatter_ptr)(mlx_array* res, const mlx_array a, const mlx_array mask, const mlx_array src, const mlx_stream s);
-extern int (*mlx_matmul_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_max_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s);
-extern int (*mlx_max_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s);
-extern int (*mlx_max_ptr)(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s);
-extern int (*mlx_maximum_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_mean_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s);
-extern int (*mlx_mean_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s);
-extern int (*mlx_mean_ptr)(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s);
-extern int (*mlx_median_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s);
-extern int (*mlx_meshgrid_ptr)(mlx_vector_array* res, const mlx_vector_array arrays, bool sparse, const char* indexing, const mlx_stream s);
-extern int (*mlx_min_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s);
-extern int (*mlx_min_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s);
-extern int (*mlx_min_ptr)(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s);
-extern int (*mlx_minimum_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_moveaxis_ptr)(mlx_array* res, const mlx_array a, int source, int destination, const mlx_stream s);
-extern int (*mlx_multiply_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_nan_to_num_ptr)(mlx_array* res, const mlx_array a, float nan, mlx_optional_float posinf, mlx_optional_float neginf, const mlx_stream s);
-extern int (*mlx_negative_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_not_equal_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_number_of_elements_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool inverted, mlx_dtype dtype, const mlx_stream s);
-extern int (*mlx_ones_ptr)(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_stream s);
-extern int (*mlx_ones_like_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_outer_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_pad_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, const int* low_pad_size, size_t low_pad_size_num, const int* high_pad_size, size_t high_pad_size_num, const mlx_array pad_value, const char* mode, const mlx_stream s);
-extern int (*mlx_pad_symmetric_ptr)(mlx_array* res, const mlx_array a, int pad_width, const mlx_array pad_value, const char* mode, const mlx_stream s);
-extern int (*mlx_partition_axis_ptr)(mlx_array* res, const mlx_array a, int kth, int axis, const mlx_stream s);
-extern int (*mlx_partition_ptr)(mlx_array* res, const mlx_array a, int kth, const mlx_stream s);
-extern int (*mlx_power_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_prod_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s);
-extern int (*mlx_prod_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s);
-extern int (*mlx_prod_ptr)(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s);
-extern int (*mlx_put_along_axis_ptr)(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array values, int axis, const mlx_stream s);
-extern int (*mlx_qqmm_ptr)(mlx_array* res, const mlx_array x, const mlx_array w, const mlx_array w_scales , mlx_optional_int group_size, mlx_optional_int bits, const char* mode, const mlx_stream s);
-extern int (*mlx_quantize_ptr)(mlx_vector_array* res, const mlx_array w, mlx_optional_int group_size, mlx_optional_int bits, const char* mode, const mlx_stream s);
-extern int (*mlx_quantized_matmul_ptr)(mlx_array* res, const mlx_array x, const mlx_array w, const mlx_array scales, const mlx_array biases , bool transpose, mlx_optional_int group_size, mlx_optional_int bits, const char* mode, const mlx_stream s);
-extern int (*mlx_radians_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_real_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_reciprocal_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_remainder_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_repeat_axis_ptr)(mlx_array* res, const mlx_array arr, int repeats, int axis, const mlx_stream s);
-extern int (*mlx_repeat_ptr)(mlx_array* res, const mlx_array arr, int repeats, const mlx_stream s);
-extern int (*mlx_reshape_ptr)(mlx_array* res, const mlx_array a, const int* shape, size_t shape_num, const mlx_stream s);
-extern int (*mlx_right_shift_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_roll_axis_ptr)(mlx_array* res, const mlx_array a, const int* shift, size_t shift_num, int axis, const mlx_stream s);
-extern int (*mlx_roll_axes_ptr)(mlx_array* res, const mlx_array a, const int* shift, size_t shift_num, const int* axes, size_t axes_num, const mlx_stream s);
-extern int (*mlx_roll_ptr)(mlx_array* res, const mlx_array a, const int* shift, size_t shift_num, const mlx_stream s);
-extern int (*mlx_round_ptr)(mlx_array* res, const mlx_array a, int decimals, const mlx_stream s);
-extern int (*mlx_rsqrt_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_scatter_ptr)(mlx_array* res, const mlx_array a, const mlx_vector_array indices, const mlx_array updates, const int* axes, size_t axes_num, const mlx_stream s);
-extern int (*mlx_scatter_single_ptr)(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array updates, int axis, const mlx_stream s);
-extern int (*mlx_scatter_add_ptr)(mlx_array* res, const mlx_array a, const mlx_vector_array indices, const mlx_array updates, const int* axes, size_t axes_num, const mlx_stream s);
-extern int (*mlx_scatter_add_single_ptr)(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array updates, int axis, const mlx_stream s);
-extern int (*mlx_scatter_add_axis_ptr)(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array values, int axis, const mlx_stream s);
-extern int (*mlx_scatter_max_ptr)(mlx_array* res, const mlx_array a, const mlx_vector_array indices, const mlx_array updates, const int* axes, size_t axes_num, const mlx_stream s);
-extern int (*mlx_scatter_max_single_ptr)(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array updates, int axis, const mlx_stream s);
-extern int (*mlx_scatter_min_ptr)(mlx_array* res, const mlx_array a, const mlx_vector_array indices, const mlx_array updates, const int* axes, size_t axes_num, const mlx_stream s);
-extern int (*mlx_scatter_min_single_ptr)(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array updates, int axis, const mlx_stream s);
-extern int (*mlx_scatter_prod_ptr)(mlx_array* res, const mlx_array a, const mlx_vector_array indices, const mlx_array updates, const int* axes, size_t axes_num, const mlx_stream s);
-extern int (*mlx_scatter_prod_single_ptr)(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array updates, int axis, const mlx_stream s);
-extern int (*mlx_segmented_mm_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_array segments, const mlx_stream s);
-extern int (*mlx_sigmoid_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_sign_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_sin_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_sinh_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_slice_ptr)(mlx_array* res, const mlx_array a, const int* start, size_t start_num, const int* stop, size_t stop_num, const int* strides, size_t strides_num, const mlx_stream s);
-extern int (*mlx_slice_dynamic_ptr)(mlx_array* res, const mlx_array a, const mlx_array start, const int* axes, size_t axes_num, const int* slice_size, size_t slice_size_num, const mlx_stream s);
-extern int (*mlx_slice_update_ptr)(mlx_array* res, const mlx_array src, const mlx_array update, const int* start, size_t start_num, const int* stop, size_t stop_num, const int* strides, size_t strides_num, const mlx_stream s);
-extern int (*mlx_slice_update_dynamic_ptr)(mlx_array* res, const mlx_array src, const mlx_array update, const mlx_array start, const int* axes, size_t axes_num, const mlx_stream s);
-extern int (*mlx_softmax_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool precise, const mlx_stream s);
-extern int (*mlx_softmax_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool precise, const mlx_stream s);
-extern int (*mlx_softmax_ptr)(mlx_array* res, const mlx_array a, bool precise, const mlx_stream s);
-extern int (*mlx_sort_axis_ptr)(mlx_array* res, const mlx_array a, int axis, const mlx_stream s);
-extern int (*mlx_sort_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_split_ptr)(mlx_vector_array* res, const mlx_array a, int num_splits, int axis, const mlx_stream s);
-extern int (*mlx_split_sections_ptr)(mlx_vector_array* res, const mlx_array a, const int* indices, size_t indices_num, int axis, const mlx_stream s);
-extern int (*mlx_sqrt_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_square_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_squeeze_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, const mlx_stream s);
-extern int (*mlx_squeeze_axis_ptr)(mlx_array* res, const mlx_array a, int axis, const mlx_stream s);
-extern int (*mlx_squeeze_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_stack_axis_ptr)(mlx_array* res, const mlx_vector_array arrays, int axis, const mlx_stream s);
-extern int (*mlx_stack_ptr)(mlx_array* res, const mlx_vector_array arrays, const mlx_stream s);
-extern int (*mlx_std_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, int ddof, const mlx_stream s);
-extern int (*mlx_std_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool keepdims, int ddof, const mlx_stream s);
-extern int (*mlx_std_ptr)(mlx_array* res, const mlx_array a, bool keepdims, int ddof, const mlx_stream s);
-extern int (*mlx_stop_gradient_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_subtract_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-extern int (*mlx_sum_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s);
-extern int (*mlx_sum_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s);
-extern int (*mlx_sum_ptr)(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s);
-extern int (*mlx_swapaxes_ptr)(mlx_array* res, const mlx_array a, int axis1, int axis2, const mlx_stream s);
-extern int (*mlx_take_axis_ptr)(mlx_array* res, const mlx_array a, const mlx_array indices, int axis, const mlx_stream s);
-extern int (*mlx_take_ptr)(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_stream s);
-extern int (*mlx_take_along_axis_ptr)(mlx_array* res, const mlx_array a, const mlx_array indices, int axis, const mlx_stream s);
-extern int (*mlx_tan_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_tanh_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_tensordot_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, const int* axes_a, size_t axes_a_num, const int* axes_b, size_t axes_b_num, const mlx_stream s);
-extern int (*mlx_tensordot_axis_ptr)(mlx_array* res, const mlx_array a, const mlx_array b, int axis, const mlx_stream s);
-extern int (*mlx_tile_ptr)(mlx_array* res, const mlx_array arr, const int* reps, size_t reps_num, const mlx_stream s);
-extern int (*mlx_to_fp8_ptr)(mlx_array* res, const mlx_array x, const mlx_stream s);
-extern int (*mlx_topk_axis_ptr)(mlx_array* res, const mlx_array a, int k, int axis, const mlx_stream s);
-extern int (*mlx_topk_ptr)(mlx_array* res, const mlx_array a, int k, const mlx_stream s);
-extern int (*mlx_trace_ptr)(mlx_array* res, const mlx_array a, int offset, int axis1, int axis2, mlx_dtype dtype, const mlx_stream s);
-extern int (*mlx_transpose_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, const mlx_stream s);
-extern int (*mlx_transpose_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_tri_ptr)(mlx_array* res, int n, int m, int k, mlx_dtype type, const mlx_stream s);
-extern int (*mlx_tril_ptr)(mlx_array* res, const mlx_array x, int k, const mlx_stream s);
-extern int (*mlx_triu_ptr)(mlx_array* res, const mlx_array x, int k, const mlx_stream s);
-extern int (*mlx_unflatten_ptr)(mlx_array* res, const mlx_array a, int axis, const int* shape, size_t shape_num, const mlx_stream s);
-extern int (*mlx_var_axes_ptr)(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, int ddof, const mlx_stream s);
-extern int (*mlx_var_axis_ptr)(mlx_array* res, const mlx_array a, int axis, bool keepdims, int ddof, const mlx_stream s);
-extern int (*mlx_var_ptr)(mlx_array* res, const mlx_array a, bool keepdims, int ddof, const mlx_stream s);
-extern int (*mlx_view_ptr)(mlx_array* res, const mlx_array a, mlx_dtype dtype, const mlx_stream s);
-extern int (*mlx_where_ptr)(mlx_array* res, const mlx_array condition, const mlx_array x, const mlx_array y, const mlx_stream s);
-extern int (*mlx_zeros_ptr)(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_stream s);
-extern int (*mlx_zeros_like_ptr)(mlx_array* res, const mlx_array a, const mlx_stream s);
-extern int (*mlx_random_bernoulli_ptr)(mlx_array* res, const mlx_array p, const int* shape, size_t shape_num, const mlx_array key , const mlx_stream s);
-extern int (*mlx_random_bits_ptr)(mlx_array* res, const int* shape, size_t shape_num, int width, const mlx_array key , const mlx_stream s);
-extern int (*mlx_random_categorical_shape_ptr)(mlx_array* res, const mlx_array logits, int axis, const int* shape, size_t shape_num, const mlx_array key , const mlx_stream s);
-extern int (*mlx_random_categorical_num_samples_ptr)(mlx_array* res, const mlx_array logits_, int axis, int num_samples, const mlx_array key , const mlx_stream s);
-extern int (*mlx_random_categorical_ptr)(mlx_array* res, const mlx_array logits, int axis, const mlx_array key , const mlx_stream s);
-extern int (*mlx_random_gumbel_ptr)(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_array key , const mlx_stream s);
-extern int (*mlx_random_key_ptr)(mlx_array* res, uint64_t seed);
-extern int (*mlx_random_laplace_ptr)(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, float loc, float scale, const mlx_array key , const mlx_stream s);
-extern int (*mlx_random_multivariate_normal_ptr)(mlx_array* res, const mlx_array mean, const mlx_array cov, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_array key , const mlx_stream s);
-extern int (*mlx_random_normal_broadcast_ptr)(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_array loc , const mlx_array scale , const mlx_array key , const mlx_stream s);
-extern int (*mlx_random_normal_ptr)(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, float loc, float scale, const mlx_array key , const mlx_stream s);
-extern int (*mlx_random_permutation_ptr)(mlx_array* res, const mlx_array x, int axis, const mlx_array key , const mlx_stream s);
-extern int (*mlx_random_permutation_arange_ptr)(mlx_array* res, int x, const mlx_array key , const mlx_stream s);
-extern int (*mlx_random_randint_ptr)(mlx_array* res, const mlx_array low, const mlx_array high, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_array key , const mlx_stream s);
-extern int (*mlx_random_seed_ptr)(uint64_t seed);
-extern int (*mlx_random_split_num_ptr)(mlx_array* res, const mlx_array key, int num, const mlx_stream s);
-extern int (*mlx_random_split_ptr)(mlx_array* res_0, mlx_array* res_1, const mlx_array key, const mlx_stream s);
-extern int (*mlx_random_truncated_normal_ptr)(mlx_array* res, const mlx_array lower, const mlx_array upper, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_array key , const mlx_stream s);
-extern int (*mlx_random_uniform_ptr)(mlx_array* res, const mlx_array low, const mlx_array high, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_array key , const mlx_stream s);
-extern mlx_stream (*mlx_stream_new_ptr)(void);
-extern mlx_stream (*mlx_stream_new_device_ptr)(mlx_device dev);
-extern int (*mlx_stream_set_ptr)(mlx_stream* stream, const mlx_stream src);
-extern int (*mlx_stream_free_ptr)(mlx_stream stream);
-extern int (*mlx_stream_tostring_ptr)(mlx_string* str, mlx_stream stream);
-extern bool (*mlx_stream_equal_ptr)(mlx_stream lhs, mlx_stream rhs);
-extern int (*mlx_stream_get_device_ptr)(mlx_device* dev, mlx_stream stream);
-extern int (*mlx_stream_get_index_ptr)(int* index, mlx_stream stream);
-extern int (*mlx_synchronize_ptr)(mlx_stream stream);
-extern int (*mlx_get_default_stream_ptr)(mlx_stream* stream, mlx_device dev);
-extern int (*mlx_set_default_stream_ptr)(mlx_stream stream);
-extern mlx_stream (*mlx_default_cpu_stream_new_ptr)(void);
-extern mlx_stream (*mlx_default_gpu_stream_new_ptr)(void);
-extern mlx_string (*mlx_string_new_ptr)(void);
-extern mlx_string (*mlx_string_new_data_ptr)(const char* str);
-extern int (*mlx_string_set_ptr)(mlx_string* str, const mlx_string src);
-extern const char* (*mlx_string_data_ptr)(mlx_string str);
-extern int (*mlx_string_free_ptr)(mlx_string str);
-extern int (*mlx_async_eval_ptr)(const mlx_vector_array outputs);
-extern int (*mlx_checkpoint_ptr)(mlx_closure* res, const mlx_closure fun);
-extern int (*mlx_custom_function_ptr)(mlx_closure* res, const mlx_closure fun, const mlx_closure_custom fun_vjp , const mlx_closure_custom_jvp fun_jvp , const mlx_closure_custom_vmap fun_vmap);
-extern int (*mlx_custom_vjp_ptr)(mlx_closure* res, const mlx_closure fun, const mlx_closure_custom fun_vjp);
-extern int (*mlx_eval_ptr)(const mlx_vector_array outputs);
-extern int (*mlx_jvp_ptr)(mlx_vector_array* res_0, mlx_vector_array* res_1, const mlx_closure fun, const mlx_vector_array primals, const mlx_vector_array tangents);
-extern int (*mlx_value_and_grad_ptr)(mlx_closure_value_and_grad* res, const mlx_closure fun, const int* argnums, size_t argnums_num);
-extern int (*mlx_vjp_ptr)(mlx_vector_array* res_0, mlx_vector_array* res_1, const mlx_closure fun, const mlx_vector_array primals, const mlx_vector_array cotangents);
-extern int (*mlx_detail_vmap_replace_ptr)(mlx_vector_array* res, const mlx_vector_array inputs, const mlx_vector_array s_inputs, const mlx_vector_array s_outputs, const int* in_axes, size_t in_axes_num, const int* out_axes, size_t out_axes_num);
-extern int (*mlx_detail_vmap_trace_ptr)(mlx_vector_array* res_0, mlx_vector_array* res_1, const mlx_closure fun, const mlx_vector_array inputs, const int* in_axes, size_t in_axes_num);
-extern mlx_vector_array (*mlx_vector_array_new_ptr)(void);
-extern int (*mlx_vector_array_set_ptr)(mlx_vector_array* vec, const mlx_vector_array src);
-extern int (*mlx_vector_array_free_ptr)(mlx_vector_array vec);
-extern mlx_vector_array (*mlx_vector_array_new_data_ptr)(const mlx_array* data, size_t size);
-extern mlx_vector_array (*mlx_vector_array_new_value_ptr)(const mlx_array val);
-extern int (*mlx_vector_array_set_data_ptr)(mlx_vector_array* vec, const mlx_array* data, size_t size);
-extern int (*mlx_vector_array_set_value_ptr)(mlx_vector_array* vec, const mlx_array val);
-extern int (*mlx_vector_array_append_data_ptr)(mlx_vector_array vec, const mlx_array* data, size_t size);
-extern int (*mlx_vector_array_append_value_ptr)(mlx_vector_array vec, const mlx_array val);
-extern size_t (*mlx_vector_array_size_ptr)(mlx_vector_array vec);
-extern int (*mlx_vector_array_get_ptr)(mlx_array* res, const mlx_vector_array vec, size_t idx);
-extern mlx_vector_vector_array (*mlx_vector_vector_array_new_ptr)(void);
-extern int (*mlx_vector_vector_array_set_ptr)(mlx_vector_vector_array* vec, const mlx_vector_vector_array src);
-extern int (*mlx_vector_vector_array_free_ptr)(mlx_vector_vector_array vec);
-extern mlx_vector_vector_array (*mlx_vector_vector_array_new_data_ptr)(const mlx_vector_array* data, size_t size);
-extern mlx_vector_vector_array (*mlx_vector_vector_array_new_value_ptr)(const mlx_vector_array val);
-extern int (*mlx_vector_vector_array_set_data_ptr)(mlx_vector_vector_array* vec, const mlx_vector_array* data, size_t size);
-extern int (*mlx_vector_vector_array_set_value_ptr)(mlx_vector_vector_array* vec, const mlx_vector_array val);
-extern int (*mlx_vector_vector_array_append_data_ptr)(mlx_vector_vector_array vec, const mlx_vector_array* data, size_t size);
-extern int (*mlx_vector_vector_array_append_value_ptr)(mlx_vector_vector_array vec, const mlx_vector_array val);
-extern size_t (*mlx_vector_vector_array_size_ptr)(mlx_vector_vector_array vec);
-extern int (*mlx_vector_vector_array_get_ptr)(mlx_vector_array* res, const mlx_vector_vector_array vec, size_t idx);
-extern mlx_vector_int (*mlx_vector_int_new_ptr)(void);
-extern int (*mlx_vector_int_set_ptr)(mlx_vector_int* vec, const mlx_vector_int src);
-extern int (*mlx_vector_int_free_ptr)(mlx_vector_int vec);
-extern mlx_vector_int (*mlx_vector_int_new_data_ptr)(int* data, size_t size);
-extern mlx_vector_int (*mlx_vector_int_new_value_ptr)(int val);
-extern int (*mlx_vector_int_set_data_ptr)(mlx_vector_int* vec, int* data, size_t size);
-extern int (*mlx_vector_int_set_value_ptr)(mlx_vector_int* vec, int val);
-extern int (*mlx_vector_int_append_data_ptr)(mlx_vector_int vec, int* data, size_t size);
-extern int (*mlx_vector_int_append_value_ptr)(mlx_vector_int vec, int val);
-extern size_t (*mlx_vector_int_size_ptr)(mlx_vector_int vec);
-extern int (*mlx_vector_int_get_ptr)(int* res, const mlx_vector_int vec, size_t idx);
-extern mlx_vector_string (*mlx_vector_string_new_ptr)(void);
-extern int (*mlx_vector_string_set_ptr)(mlx_vector_string* vec, const mlx_vector_string src);
-extern int (*mlx_vector_string_free_ptr)(mlx_vector_string vec);
-extern mlx_vector_string (*mlx_vector_string_new_data_ptr)(const char** data, size_t size);
-extern mlx_vector_string (*mlx_vector_string_new_value_ptr)(const char* val);
-extern int (*mlx_vector_string_set_data_ptr)(mlx_vector_string* vec, const char** data, size_t size);
-extern int (*mlx_vector_string_set_value_ptr)(mlx_vector_string* vec, const char* val);
-extern int (*mlx_vector_string_append_data_ptr)(mlx_vector_string vec, const char** data, size_t size);
-extern int (*mlx_vector_string_append_value_ptr)(mlx_vector_string vec, const char* val);
-extern size_t (*mlx_vector_string_size_ptr)(mlx_vector_string vec);
-extern int (*mlx_vector_string_get_ptr)(char** res, const mlx_vector_string vec, size_t idx);
-extern int (*mlx_version_ptr)(mlx_string* str_);
-
-// Initialize all function pointers via dlsym (defined in mlx.c)
-int mlx_load_functions(void* handle);
-
-// Wrapper function declarations that call through function pointers
-// Go code calls these directly as C.mlx_* (no #define redirection needed)
-size_t mlx_dtype_size(mlx_dtype dtype);
-
-int mlx_array_tostring(mlx_string* str, const mlx_array arr);
-
-mlx_array mlx_array_new(void);
-
-int mlx_array_free(mlx_array arr);
-
-mlx_array mlx_array_new_bool(bool val);
-
-mlx_array mlx_array_new_int(int val);
-
-mlx_array mlx_array_new_float32(float val);
-
-mlx_array mlx_array_new_float(float val);
-
-mlx_array mlx_array_new_float64(double val);
-
-mlx_array mlx_array_new_double(double val);
-
-mlx_array mlx_array_new_complex(float real_val, float imag_val);
-
-mlx_array mlx_array_new_data(const void* data, const int* shape, int dim, mlx_dtype dtype);
-
-int mlx_array_set(mlx_array* arr, const mlx_array src);
-
-int mlx_array_set_bool(mlx_array* arr, bool val);
-
-int mlx_array_set_int(mlx_array* arr, int val);
-
-int mlx_array_set_float32(mlx_array* arr, float val);
-
-int mlx_array_set_float(mlx_array* arr, float val);
-
-int mlx_array_set_float64(mlx_array* arr, double val);
-
-int mlx_array_set_double(mlx_array* arr, double val);
-
-int mlx_array_set_complex(mlx_array* arr, float real_val, float imag_val);
-
-int mlx_array_set_data(mlx_array* arr, const void* data, const int* shape, int dim, mlx_dtype dtype);
-
-size_t mlx_array_itemsize(const mlx_array arr);
-
-size_t mlx_array_size(const mlx_array arr);
-
-size_t mlx_array_nbytes(const mlx_array arr);
-
-size_t mlx_array_ndim(const mlx_array arr);
-
-const int* mlx_array_shape(const mlx_array arr);
-
-const size_t* mlx_array_strides(const mlx_array arr);
-
-int mlx_array_dim(const mlx_array arr, int dim);
-
-mlx_dtype mlx_array_dtype(const mlx_array arr);
-
-int mlx_array_eval(mlx_array arr);
-
-int mlx_array_item_bool(bool* res, const mlx_array arr);
-
-int mlx_array_item_uint8(uint8_t* res, const mlx_array arr);
-
-int mlx_array_item_uint16(uint16_t* res, const mlx_array arr);
-
-int mlx_array_item_uint32(uint32_t* res, const mlx_array arr);
-
-int mlx_array_item_uint64(uint64_t* res, const mlx_array arr);
-
-int mlx_array_item_int8(int8_t* res, const mlx_array arr);
-
-int mlx_array_item_int16(int16_t* res, const mlx_array arr);
-
-int mlx_array_item_int32(int32_t* res, const mlx_array arr);
-
-int mlx_array_item_int64(int64_t* res, const mlx_array arr);
-
-int mlx_array_item_float32(float* res, const mlx_array arr);
-
-int mlx_array_item_float64(double* res, const mlx_array arr);
-
-int mlx_array_item_complex64(float _Complex* res, const mlx_array arr);
-
-#if defined(__aarch64__) || defined(_M_ARM64)
-int mlx_array_item_float16(float16_t* res, const mlx_array arr);
-#endif
-
-#if defined(__aarch64__) || defined(_M_ARM64)
-int mlx_array_item_bfloat16(bfloat16_t* res, const mlx_array arr);
-#endif
-
-const bool* mlx_array_data_bool(const mlx_array arr);
-
-const uint8_t* mlx_array_data_uint8(const mlx_array arr);
-
-const uint16_t* mlx_array_data_uint16(const mlx_array arr);
-
-const uint32_t* mlx_array_data_uint32(const mlx_array arr);
-
-const uint64_t* mlx_array_data_uint64(const mlx_array arr);
-
-const int8_t* mlx_array_data_int8(const mlx_array arr);
-
-const int16_t* mlx_array_data_int16(const mlx_array arr);
-
-const int32_t* mlx_array_data_int32(const mlx_array arr);
-
-const int64_t* mlx_array_data_int64(const mlx_array arr);
-
-const float* mlx_array_data_float32(const mlx_array arr);
-
-const double* mlx_array_data_float64(const mlx_array arr);
-
-const float _Complex* mlx_array_data_complex64(const mlx_array arr);
-
-#if defined(__aarch64__) || defined(_M_ARM64)
-const float16_t* mlx_array_data_float16(const mlx_array arr);
-#endif
-
-#if defined(__aarch64__) || defined(_M_ARM64)
-const bfloat16_t* mlx_array_data_bfloat16(const mlx_array arr);
-#endif
-
-int _mlx_array_is_available(bool* res, const mlx_array arr);
-
-int _mlx_array_wait(const mlx_array arr);
-
-int _mlx_array_is_contiguous(bool* res, const mlx_array arr);
-
-int _mlx_array_is_row_contiguous(bool* res, const mlx_array arr);
-
-int _mlx_array_is_col_contiguous(bool* res, const mlx_array arr);
-
-mlx_closure mlx_closure_new(void);
-
-int mlx_closure_free(mlx_closure cls);
-
-mlx_closure mlx_closure_new_func(int (*fun)(mlx_vector_array*, const mlx_vector_array));
-
-mlx_closure mlx_closure_new_func_payload(int (*fun)(mlx_vector_array*, const mlx_vector_array, void*), void* payload, void (*dtor)(void*));
-
-int mlx_closure_set(mlx_closure* cls, const mlx_closure src);
-
-int mlx_closure_apply(mlx_vector_array* res, mlx_closure cls, const mlx_vector_array input);
-
-mlx_closure mlx_closure_new_unary(int (*fun)(mlx_array*, const mlx_array));
-
-mlx_closure_kwargs mlx_closure_kwargs_new(void);
-
-int mlx_closure_kwargs_free(mlx_closure_kwargs cls);
-
-mlx_closure_kwargs mlx_closure_kwargs_new_func(int (*fun)( mlx_vector_array*, const mlx_vector_array, const mlx_map_string_to_array));
-
-mlx_closure_kwargs mlx_closure_kwargs_new_func_payload(int (*fun)( mlx_vector_array*, const mlx_vector_array, const mlx_map_string_to_array, void*), void* payload, void (*dtor)(void*));
-
-int mlx_closure_kwargs_set(mlx_closure_kwargs* cls, const mlx_closure_kwargs src);
-
-int mlx_closure_kwargs_apply(mlx_vector_array* res, mlx_closure_kwargs cls, const mlx_vector_array input_0, const mlx_map_string_to_array input_1);
-
-mlx_closure_value_and_grad mlx_closure_value_and_grad_new(void);
-
-int mlx_closure_value_and_grad_free(mlx_closure_value_and_grad cls);
-
-mlx_closure_value_and_grad mlx_closure_value_and_grad_new_func(int (*fun)(mlx_vector_array*, mlx_vector_array*, const mlx_vector_array));
-
-mlx_closure_value_and_grad mlx_closure_value_and_grad_new_func_payload(int (*fun)( mlx_vector_array*, mlx_vector_array*, const mlx_vector_array, void*), void* payload, void (*dtor)(void*));
-
-int mlx_closure_value_and_grad_set(mlx_closure_value_and_grad* cls, const mlx_closure_value_and_grad src);
-
-int mlx_closure_value_and_grad_apply(mlx_vector_array* res_0, mlx_vector_array* res_1, mlx_closure_value_and_grad cls, const mlx_vector_array input);
-
-mlx_closure_custom mlx_closure_custom_new(void);
-
-int mlx_closure_custom_free(mlx_closure_custom cls);
-
-mlx_closure_custom mlx_closure_custom_new_func(int (*fun)( mlx_vector_array*, const mlx_vector_array, const mlx_vector_array, const mlx_vector_array));
-
-mlx_closure_custom mlx_closure_custom_new_func_payload(int (*fun)( mlx_vector_array*, const mlx_vector_array, const mlx_vector_array, const mlx_vector_array, void*), void* payload, void (*dtor)(void*));
-
-int mlx_closure_custom_set(mlx_closure_custom* cls, const mlx_closure_custom src);
-
-int mlx_closure_custom_apply(mlx_vector_array* res, mlx_closure_custom cls, const mlx_vector_array input_0, const mlx_vector_array input_1, const mlx_vector_array input_2);
-
-mlx_closure_custom_jvp mlx_closure_custom_jvp_new(void);
-
-int mlx_closure_custom_jvp_free(mlx_closure_custom_jvp cls);
-
-mlx_closure_custom_jvp mlx_closure_custom_jvp_new_func(int (*fun)( mlx_vector_array*, const mlx_vector_array, const mlx_vector_array, const int*, size_t _num));
-
-mlx_closure_custom_jvp mlx_closure_custom_jvp_new_func_payload(int (*fun)( mlx_vector_array*, const mlx_vector_array, const mlx_vector_array, const int*, size_t _num, void*), void* payload, void (*dtor)(void*));
-
-int mlx_closure_custom_jvp_set(mlx_closure_custom_jvp* cls, const mlx_closure_custom_jvp src);
-
-int mlx_closure_custom_jvp_apply(mlx_vector_array* res, mlx_closure_custom_jvp cls, const mlx_vector_array input_0, const mlx_vector_array input_1, const int* input_2, size_t input_2_num);
-
-mlx_closure_custom_vmap mlx_closure_custom_vmap_new(void);
-
-int mlx_closure_custom_vmap_free(mlx_closure_custom_vmap cls);
-
-mlx_closure_custom_vmap mlx_closure_custom_vmap_new_func(int (*fun)( mlx_vector_array*, mlx_vector_int*, const mlx_vector_array, const int*, size_t _num));
-
-mlx_closure_custom_vmap mlx_closure_custom_vmap_new_func_payload(int (*fun)( mlx_vector_array*, mlx_vector_int*, const mlx_vector_array, const int*, size_t _num, void*), void* payload, void (*dtor)(void*));
-
-int mlx_closure_custom_vmap_set(mlx_closure_custom_vmap* cls, const mlx_closure_custom_vmap src);
-
-int mlx_closure_custom_vmap_apply(mlx_vector_array* res_0, mlx_vector_int* res_1, mlx_closure_custom_vmap cls, const mlx_vector_array input_0, const int* input_1, size_t input_1_num);
-
-int mlx_compile(mlx_closure* res, const mlx_closure fun, bool shapeless);
-
-int mlx_detail_compile(mlx_closure* res, const mlx_closure fun, uintptr_t fun_id, bool shapeless, const uint64_t* constants, size_t constants_num);
-
-int mlx_detail_compile_clear_cache(void);
-
-int mlx_detail_compile_erase(uintptr_t fun_id);
-
-int mlx_disable_compile(void);
-
-int mlx_enable_compile(void);
-
-int mlx_set_compile_mode(mlx_compile_mode mode);
-
-mlx_device mlx_device_new(void);
-
-mlx_device mlx_device_new_type(mlx_device_type type, int index);
-
-int mlx_device_free(mlx_device dev);
-
-int mlx_device_set(mlx_device* dev, const mlx_device src);
-
-int mlx_device_tostring(mlx_string* str, mlx_device dev);
-
-bool mlx_device_equal(mlx_device lhs, mlx_device rhs);
-
-int mlx_device_get_index(int* index, mlx_device dev);
-
-int mlx_device_get_type(mlx_device_type* type, mlx_device dev);
-
-int mlx_get_default_device(mlx_device* dev);
-
-int mlx_set_default_device(mlx_device dev);
-
-int mlx_distributed_all_gather(mlx_array* res, const mlx_array x, const mlx_distributed_group group , const mlx_stream S);
-
-int mlx_distributed_all_max(mlx_array* res, const mlx_array x, const mlx_distributed_group group , const mlx_stream s);
-
-int mlx_distributed_all_min(mlx_array* res, const mlx_array x, const mlx_distributed_group group , const mlx_stream s);
-
-int mlx_distributed_all_sum(mlx_array* res, const mlx_array x, const mlx_distributed_group group , const mlx_stream s);
-
-int mlx_distributed_recv(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, int src, const mlx_distributed_group group , const mlx_stream s);
-
-int mlx_distributed_recv_like(mlx_array* res, const mlx_array x, int src, const mlx_distributed_group group , const mlx_stream s);
-
-int mlx_distributed_send(mlx_array* res, const mlx_array x, int dst, const mlx_distributed_group group , const mlx_stream s);
-
-int mlx_distributed_sum_scatter(mlx_array* res, const mlx_array x, const mlx_distributed_group group , const mlx_stream s);
-
-int mlx_distributed_group_rank(mlx_distributed_group group);
-
-int mlx_distributed_group_size(mlx_distributed_group group);
-
-mlx_distributed_group mlx_distributed_group_split(mlx_distributed_group group, int color, int key);
-
-bool mlx_distributed_is_available(void);
-
-mlx_distributed_group mlx_distributed_init(bool strict);
-
-void mlx_set_error_handler(mlx_error_handler_func handler, void* data, void (*dtor)(void*));
-
-void _mlx_error(const char* file, const int line, const char* fmt, ...);
-
-int mlx_export_function(const char* file, const mlx_closure fun, const mlx_vector_array args, bool shapeless);
-
-int mlx_export_function_kwargs(const char* file, const mlx_closure_kwargs fun, const mlx_vector_array args, const mlx_map_string_to_array kwargs, bool shapeless);
-
-mlx_function_exporter mlx_function_exporter_new(const char* file, const mlx_closure fun, bool shapeless);
-
-int mlx_function_exporter_free(mlx_function_exporter xfunc);
-
-int mlx_function_exporter_apply(const mlx_function_exporter xfunc, const mlx_vector_array args);
-
-int mlx_function_exporter_apply_kwargs(const mlx_function_exporter xfunc, const mlx_vector_array args, const mlx_map_string_to_array kwargs);
-
-mlx_imported_function mlx_imported_function_new(const char* file);
-
-int mlx_imported_function_free(mlx_imported_function xfunc);
-
-int mlx_imported_function_apply(mlx_vector_array* res, const mlx_imported_function xfunc, const mlx_vector_array args);
-
-int mlx_imported_function_apply_kwargs(mlx_vector_array* res, const mlx_imported_function xfunc, const mlx_vector_array args, const mlx_map_string_to_array kwargs);
-
-mlx_fast_cuda_kernel_config mlx_fast_cuda_kernel_config_new(void);
-
-void mlx_fast_cuda_kernel_config_free(mlx_fast_cuda_kernel_config cls);
-
-int mlx_fast_cuda_kernel_config_add_output_arg(mlx_fast_cuda_kernel_config cls, const int* shape, size_t size, mlx_dtype dtype);
-
-int mlx_fast_cuda_kernel_config_set_grid(mlx_fast_cuda_kernel_config cls, int grid1, int grid2, int grid3);
-
-int mlx_fast_cuda_kernel_config_set_thread_group(mlx_fast_cuda_kernel_config cls, int thread1, int thread2, int thread3);
-
-int mlx_fast_cuda_kernel_config_set_init_value(mlx_fast_cuda_kernel_config cls, float value);
-
-int mlx_fast_cuda_kernel_config_set_verbose(mlx_fast_cuda_kernel_config cls, bool verbose);
-
-int mlx_fast_cuda_kernel_config_add_template_arg_dtype(mlx_fast_cuda_kernel_config cls, const char* name, mlx_dtype dtype);
-
-int mlx_fast_cuda_kernel_config_add_template_arg_int(mlx_fast_cuda_kernel_config cls, const char* name, int value);
-
-int mlx_fast_cuda_kernel_config_add_template_arg_bool(mlx_fast_cuda_kernel_config cls, const char* name, bool value);
-
-mlx_fast_cuda_kernel mlx_fast_cuda_kernel_new(const char* name, const mlx_vector_string input_names, const mlx_vector_string output_names, const char* source, const char* header, bool ensure_row_contiguous, int shared_memory);
-
-void mlx_fast_cuda_kernel_free(mlx_fast_cuda_kernel cls);
-
-int mlx_fast_cuda_kernel_apply(mlx_vector_array* outputs, mlx_fast_cuda_kernel cls, const mlx_vector_array inputs, const mlx_fast_cuda_kernel_config config, const mlx_stream stream);
-
-int mlx_fast_layer_norm(mlx_array* res, const mlx_array x, const mlx_array weight , const mlx_array bias , float eps, const mlx_stream s);
-
-mlx_fast_metal_kernel_config mlx_fast_metal_kernel_config_new(void);
-
-void mlx_fast_metal_kernel_config_free(mlx_fast_metal_kernel_config cls);
-
-int mlx_fast_metal_kernel_config_add_output_arg(mlx_fast_metal_kernel_config cls, const int* shape, size_t size, mlx_dtype dtype);
-
-int mlx_fast_metal_kernel_config_set_grid(mlx_fast_metal_kernel_config cls, int grid1, int grid2, int grid3);
-
-int mlx_fast_metal_kernel_config_set_thread_group(mlx_fast_metal_kernel_config cls, int thread1, int thread2, int thread3);
-
-int mlx_fast_metal_kernel_config_set_init_value(mlx_fast_metal_kernel_config cls, float value);
-
-int mlx_fast_metal_kernel_config_set_verbose(mlx_fast_metal_kernel_config cls, bool verbose);
-
-int mlx_fast_metal_kernel_config_add_template_arg_dtype(mlx_fast_metal_kernel_config cls, const char* name, mlx_dtype dtype);
-
-int mlx_fast_metal_kernel_config_add_template_arg_int(mlx_fast_metal_kernel_config cls, const char* name, int value);
-
-int mlx_fast_metal_kernel_config_add_template_arg_bool(mlx_fast_metal_kernel_config cls, const char* name, bool value);
-
-mlx_fast_metal_kernel mlx_fast_metal_kernel_new(const char* name, const mlx_vector_string input_names, const mlx_vector_string output_names, const char* source, const char* header, bool ensure_row_contiguous, bool atomic_outputs);
-
-void mlx_fast_metal_kernel_free(mlx_fast_metal_kernel cls);
-
-int mlx_fast_metal_kernel_apply(mlx_vector_array* outputs, mlx_fast_metal_kernel cls, const mlx_vector_array inputs, const mlx_fast_metal_kernel_config config, const mlx_stream stream);
-
-int mlx_fast_rms_norm(mlx_array* res, const mlx_array x, const mlx_array weight , float eps, const mlx_stream s);
-
-int mlx_fast_rope(mlx_array* res, const mlx_array x, int dims, bool traditional, mlx_optional_float base, float scale, int offset, const mlx_array freqs , const mlx_stream s);
-
-int mlx_fast_rope_dynamic(mlx_array* res, const mlx_array x, int dims, bool traditional, mlx_optional_float base, float scale, const mlx_array offset, const mlx_array freqs , const mlx_stream s);
-
-int mlx_fast_scaled_dot_product_attention(mlx_array* res, const mlx_array queries, const mlx_array keys, const mlx_array values, float scale, const char* mask_mode, const mlx_array mask_arr , const mlx_array sinks , const mlx_stream s);
-
-int mlx_fft_fft(mlx_array* res, const mlx_array a, int n, int axis, const mlx_stream s);
-
-int mlx_fft_fft2(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s);
-
-int mlx_fft_fftn(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s);
-
-int mlx_fft_fftshift(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, const mlx_stream s);
-
-int mlx_fft_ifft(mlx_array* res, const mlx_array a, int n, int axis, const mlx_stream s);
-
-int mlx_fft_ifft2(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s);
-
-int mlx_fft_ifftn(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s);
-
-int mlx_fft_ifftshift(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, const mlx_stream s);
-
-int mlx_fft_irfft(mlx_array* res, const mlx_array a, int n, int axis, const mlx_stream s);
-
-int mlx_fft_irfft2(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s);
-
-int mlx_fft_irfftn(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s);
-
-int mlx_fft_rfft(mlx_array* res, const mlx_array a, int n, int axis, const mlx_stream s);
-
-int mlx_fft_rfft2(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s);
-
-int mlx_fft_rfftn(mlx_array* res, const mlx_array a, const int* n, size_t n_num, const int* axes, size_t axes_num, const mlx_stream s);
-
-int mlx_load_reader(mlx_array* res, mlx_io_reader in_stream, const mlx_stream s);
-
-int mlx_load(mlx_array* res, const char* file, const mlx_stream s);
-
-int mlx_load_safetensors_reader(mlx_map_string_to_array* res_0, mlx_map_string_to_string* res_1, mlx_io_reader in_stream, const mlx_stream s);
-
-int mlx_load_safetensors(mlx_map_string_to_array* res_0, mlx_map_string_to_string* res_1, const char* file, const mlx_stream s);
-
-int mlx_save_writer(mlx_io_writer out_stream, const mlx_array a);
-
-int mlx_save(const char* file, const mlx_array a);
-
-int mlx_save_safetensors_writer(mlx_io_writer in_stream, const mlx_map_string_to_array param, const mlx_map_string_to_string metadata);
-
-int mlx_save_safetensors(const char* file, const mlx_map_string_to_array param, const mlx_map_string_to_string metadata);
-
-mlx_io_reader mlx_io_reader_new(void* desc, mlx_io_vtable vtable);
-
-int mlx_io_reader_descriptor(void** desc_, mlx_io_reader io);
-
-int mlx_io_reader_tostring(mlx_string* str_, mlx_io_reader io);
-
-int mlx_io_reader_free(mlx_io_reader io);
-
-mlx_io_writer mlx_io_writer_new(void* desc, mlx_io_vtable vtable);
-
-int mlx_io_writer_descriptor(void** desc_, mlx_io_writer io);
-
-int mlx_io_writer_tostring(mlx_string* str_, mlx_io_writer io);
-
-int mlx_io_writer_free(mlx_io_writer io);
-
-int mlx_linalg_cholesky(mlx_array* res, const mlx_array a, bool upper, const mlx_stream s);
-
-int mlx_linalg_cholesky_inv(mlx_array* res, const mlx_array a, bool upper, const mlx_stream s);
-
-int mlx_linalg_cross(mlx_array* res, const mlx_array a, const mlx_array b, int axis, const mlx_stream s);
-
-int mlx_linalg_eig(mlx_array* res_0, mlx_array* res_1, const mlx_array a, const mlx_stream s);
-
-int mlx_linalg_eigh(mlx_array* res_0, mlx_array* res_1, const mlx_array a, const char* UPLO, const mlx_stream s);
-
-int mlx_linalg_eigvals(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_linalg_eigvalsh(mlx_array* res, const mlx_array a, const char* UPLO, const mlx_stream s);
-
-int mlx_linalg_inv(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_linalg_lu(mlx_vector_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_linalg_lu_factor(mlx_array* res_0, mlx_array* res_1, const mlx_array a, const mlx_stream s);
-
-int mlx_linalg_norm(mlx_array* res, const mlx_array a, double ord, const int* axis , size_t axis_num, bool keepdims, const mlx_stream s);
-
-int mlx_linalg_norm_matrix(mlx_array* res, const mlx_array a, const char* ord, const int* axis , size_t axis_num, bool keepdims, const mlx_stream s);
-
-int mlx_linalg_norm_l2(mlx_array* res, const mlx_array a, const int* axis , size_t axis_num, bool keepdims, const mlx_stream s);
-
-int mlx_linalg_pinv(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_linalg_qr(mlx_array* res_0, mlx_array* res_1, const mlx_array a, const mlx_stream s);
-
-int mlx_linalg_solve(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_linalg_solve_triangular(mlx_array* res, const mlx_array a, const mlx_array b, bool upper, const mlx_stream s);
-
-int mlx_linalg_svd(mlx_vector_array* res, const mlx_array a, bool compute_uv, const mlx_stream s);
-
-int mlx_linalg_tri_inv(mlx_array* res, const mlx_array a, bool upper, const mlx_stream s);
-
-mlx_map_string_to_array mlx_map_string_to_array_new(void);
-
-int mlx_map_string_to_array_set(mlx_map_string_to_array* map, const mlx_map_string_to_array src);
-
-int mlx_map_string_to_array_free(mlx_map_string_to_array map);
-
-int mlx_map_string_to_array_insert(mlx_map_string_to_array map, const char* key, const mlx_array value);
-
-int mlx_map_string_to_array_get(mlx_array* value, const mlx_map_string_to_array map, const char* key);
-
-mlx_map_string_to_array_iterator mlx_map_string_to_array_iterator_new(mlx_map_string_to_array map);
-
-int mlx_map_string_to_array_iterator_free(mlx_map_string_to_array_iterator it);
-
-int mlx_map_string_to_array_iterator_next(const char** key, mlx_array* value, mlx_map_string_to_array_iterator it);
-
-mlx_map_string_to_string mlx_map_string_to_string_new(void);
-
-int mlx_map_string_to_string_set(mlx_map_string_to_string* map, const mlx_map_string_to_string src);
-
-int mlx_map_string_to_string_free(mlx_map_string_to_string map);
-
-int mlx_map_string_to_string_insert(mlx_map_string_to_string map, const char* key, const char* value);
-
-int mlx_map_string_to_string_get(const char** value, const mlx_map_string_to_string map, const char* key);
-
-mlx_map_string_to_string_iterator mlx_map_string_to_string_iterator_new(mlx_map_string_to_string map);
-
-int mlx_map_string_to_string_iterator_free(mlx_map_string_to_string_iterator it);
-
-int mlx_map_string_to_string_iterator_next(const char** key, const char** value, mlx_map_string_to_string_iterator it);
-
-int mlx_clear_cache(void);
-
-int mlx_get_active_memory(size_t* res);
-
-int mlx_get_cache_memory(size_t* res);
-
-int mlx_get_memory_limit(size_t* res);
-
-int mlx_get_peak_memory(size_t* res);
-
-int mlx_reset_peak_memory(void);
-
-int mlx_set_cache_limit(size_t* res, size_t limit);
-
-int mlx_set_memory_limit(size_t* res, size_t limit);
-
-int mlx_set_wired_limit(size_t* res, size_t limit);
-
-mlx_metal_device_info_t mlx_metal_device_info(void);
-
-int mlx_metal_is_available(bool* res);
-
-int mlx_metal_start_capture(const char* path);
-
-int mlx_metal_stop_capture(void);
-
-int mlx_abs(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_add(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_addmm(mlx_array* res, const mlx_array c, const mlx_array a, const mlx_array b, float alpha, float beta, const mlx_stream s);
-
-int mlx_all_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s);
-
-int mlx_all_axis(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s);
-
-int mlx_all(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s);
-
-int mlx_allclose(mlx_array* res, const mlx_array a, const mlx_array b, double rtol, double atol, bool equal_nan, const mlx_stream s);
-
-int mlx_any_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s);
-
-int mlx_any_axis(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s);
-
-int mlx_any(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s);
-
-int mlx_arange(mlx_array* res, double start, double stop, double step, mlx_dtype dtype, const mlx_stream s);
-
-int mlx_arccos(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_arccosh(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_arcsin(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_arcsinh(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_arctan(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_arctan2(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_arctanh(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_argmax_axis(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s);
-
-int mlx_argmax(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s);
-
-int mlx_argmin_axis(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s);
-
-int mlx_argmin(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s);
-
-int mlx_argpartition_axis(mlx_array* res, const mlx_array a, int kth, int axis, const mlx_stream s);
-
-int mlx_argpartition(mlx_array* res, const mlx_array a, int kth, const mlx_stream s);
-
-int mlx_argsort_axis(mlx_array* res, const mlx_array a, int axis, const mlx_stream s);
-
-int mlx_argsort(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_array_equal(mlx_array* res, const mlx_array a, const mlx_array b, bool equal_nan, const mlx_stream s);
-
-int mlx_as_strided(mlx_array* res, const mlx_array a, const int* shape, size_t shape_num, const int64_t* strides, size_t strides_num, size_t offset, const mlx_stream s);
-
-int mlx_astype(mlx_array* res, const mlx_array a, mlx_dtype dtype, const mlx_stream s);
-
-int mlx_atleast_1d(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_atleast_2d(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_atleast_3d(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_bitwise_and(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_bitwise_invert(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_bitwise_or(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_bitwise_xor(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_block_masked_mm(mlx_array* res, const mlx_array a, const mlx_array b, int block_size, const mlx_array mask_out , const mlx_array mask_lhs , const mlx_array mask_rhs , const mlx_stream s);
-
-int mlx_broadcast_arrays(mlx_vector_array* res, const mlx_vector_array inputs, const mlx_stream s);
-
-int mlx_broadcast_to(mlx_array* res, const mlx_array a, const int* shape, size_t shape_num, const mlx_stream s);
-
-int mlx_ceil(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_clip(mlx_array* res, const mlx_array a, const mlx_array a_min , const mlx_array a_max , const mlx_stream s);
-
-int mlx_concatenate_axis(mlx_array* res, const mlx_vector_array arrays, int axis, const mlx_stream s);
-
-int mlx_concatenate(mlx_array* res, const mlx_vector_array arrays, const mlx_stream s);
-
-int mlx_conjugate(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_contiguous(mlx_array* res, const mlx_array a, bool allow_col_major, const mlx_stream s);
-
-int mlx_conv1d(mlx_array* res, const mlx_array input, const mlx_array weight, int stride, int padding, int dilation, int groups, const mlx_stream s);
-
-int mlx_conv2d(mlx_array* res, const mlx_array input, const mlx_array weight, int stride_0, int stride_1, int padding_0, int padding_1, int dilation_0, int dilation_1, int groups, const mlx_stream s);
-
-int mlx_conv3d(mlx_array* res, const mlx_array input, const mlx_array weight, int stride_0, int stride_1, int stride_2, int padding_0, int padding_1, int padding_2, int dilation_0, int dilation_1, int dilation_2, int groups, const mlx_stream s);
-
-int mlx_conv_general(mlx_array* res, const mlx_array input, const mlx_array weight, const int* stride, size_t stride_num, const int* padding_lo, size_t padding_lo_num, const int* padding_hi, size_t padding_hi_num, const int* kernel_dilation, size_t kernel_dilation_num, const int* input_dilation, size_t input_dilation_num, int groups, bool flip, const mlx_stream s);
-
-int mlx_conv_transpose1d(mlx_array* res, const mlx_array input, const mlx_array weight, int stride, int padding, int dilation, int output_padding, int groups, const mlx_stream s);
-
-int mlx_conv_transpose2d(mlx_array* res, const mlx_array input, const mlx_array weight, int stride_0, int stride_1, int padding_0, int padding_1, int dilation_0, int dilation_1, int output_padding_0, int output_padding_1, int groups, const mlx_stream s);
-
-int mlx_conv_transpose3d(mlx_array* res, const mlx_array input, const mlx_array weight, int stride_0, int stride_1, int stride_2, int padding_0, int padding_1, int padding_2, int dilation_0, int dilation_1, int dilation_2, int output_padding_0, int output_padding_1, int output_padding_2, int groups, const mlx_stream s);
-
-int mlx_copy(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_cos(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_cosh(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_cummax(mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s);
-
-int mlx_cummin(mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s);
-
-int mlx_cumprod(mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s);
-
-int mlx_cumsum(mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s);
-
-int mlx_degrees(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_depends(mlx_vector_array* res, const mlx_vector_array inputs, const mlx_vector_array dependencies);
-
-int mlx_dequantize(mlx_array* res, const mlx_array w, const mlx_array scales, const mlx_array biases , mlx_optional_int group_size, mlx_optional_int bits, const char* mode, mlx_optional_dtype dtype, const mlx_stream s);
-
-int mlx_diag(mlx_array* res, const mlx_array a, int k, const mlx_stream s);
-
-int mlx_diagonal(mlx_array* res, const mlx_array a, int offset, int axis1, int axis2, const mlx_stream s);
-
-int mlx_divide(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_divmod(mlx_vector_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_einsum(mlx_array* res, const char* subscripts, const mlx_vector_array operands, const mlx_stream s);
-
-int mlx_equal(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_erf(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_erfinv(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_exp(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_expand_dims_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, const mlx_stream s);
-
-int mlx_expand_dims(mlx_array* res, const mlx_array a, int axis, const mlx_stream s);
-
-int mlx_expm1(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_eye(mlx_array* res, int n, int m, int k, mlx_dtype dtype, const mlx_stream s);
-
-int mlx_flatten(mlx_array* res, const mlx_array a, int start_axis, int end_axis, const mlx_stream s);
-
-int mlx_floor(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_floor_divide(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_from_fp8(mlx_array* res, const mlx_array x, mlx_dtype dtype, const mlx_stream s);
-
-int mlx_full(mlx_array* res, const int* shape, size_t shape_num, const mlx_array vals, mlx_dtype dtype, const mlx_stream s);
-
-int mlx_full_like(mlx_array* res, const mlx_array a, const mlx_array vals, mlx_dtype dtype, const mlx_stream s);
-
-int mlx_gather(mlx_array* res, const mlx_array a, const mlx_vector_array indices, const int* axes, size_t axes_num, const int* slice_sizes, size_t slice_sizes_num, const mlx_stream s);
-
-int mlx_gather_single(mlx_array* res, const mlx_array a, const mlx_array indices, int axis, const int* slice_sizes, size_t slice_sizes_num, const mlx_stream s);
-
-int mlx_gather_mm(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_array lhs_indices , const mlx_array rhs_indices , bool sorted_indices, const mlx_stream s);
-
-int mlx_gather_qmm(mlx_array* res, const mlx_array x, const mlx_array w, const mlx_array scales, const mlx_array biases , const mlx_array lhs_indices , const mlx_array rhs_indices , bool transpose, mlx_optional_int group_size, mlx_optional_int bits, const char* mode, bool sorted_indices, const mlx_stream s);
-
-int mlx_greater(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_greater_equal(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_hadamard_transform(mlx_array* res, const mlx_array a, mlx_optional_float scale, const mlx_stream s);
-
-int mlx_identity(mlx_array* res, int n, mlx_dtype dtype, const mlx_stream s);
-
-int mlx_imag(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_inner(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_isclose(mlx_array* res, const mlx_array a, const mlx_array b, double rtol, double atol, bool equal_nan, const mlx_stream s);
-
-int mlx_isfinite(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_isinf(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_isnan(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_isneginf(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_isposinf(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_kron(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_left_shift(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_less(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_less_equal(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_linspace(mlx_array* res, double start, double stop, int num, mlx_dtype dtype, const mlx_stream s);
-
-int mlx_log(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_log10(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_log1p(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_log2(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_logaddexp(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_logcumsumexp(mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s);
-
-int mlx_logical_and(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_logical_not(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_logical_or(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_logsumexp_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s);
-
-int mlx_logsumexp_axis(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s);
-
-int mlx_logsumexp(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s);
-
-int mlx_masked_scatter(mlx_array* res, const mlx_array a, const mlx_array mask, const mlx_array src, const mlx_stream s);
-
-int mlx_matmul(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_max_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s);
-
-int mlx_max_axis(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s);
-
-int mlx_max(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s);
-
-int mlx_maximum(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_mean_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s);
-
-int mlx_mean_axis(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s);
-
-int mlx_mean(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s);
-
-int mlx_median(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s);
-
-int mlx_meshgrid(mlx_vector_array* res, const mlx_vector_array arrays, bool sparse, const char* indexing, const mlx_stream s);
-
-int mlx_min_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s);
-
-int mlx_min_axis(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s);
-
-int mlx_min(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s);
-
-int mlx_minimum(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_moveaxis(mlx_array* res, const mlx_array a, int source, int destination, const mlx_stream s);
-
-int mlx_multiply(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_nan_to_num(mlx_array* res, const mlx_array a, float nan, mlx_optional_float posinf, mlx_optional_float neginf, const mlx_stream s);
-
-int mlx_negative(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_not_equal(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_number_of_elements(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool inverted, mlx_dtype dtype, const mlx_stream s);
-
-int mlx_ones(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_stream s);
-
-int mlx_ones_like(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_outer(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_pad(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, const int* low_pad_size, size_t low_pad_size_num, const int* high_pad_size, size_t high_pad_size_num, const mlx_array pad_value, const char* mode, const mlx_stream s);
-
-int mlx_pad_symmetric(mlx_array* res, const mlx_array a, int pad_width, const mlx_array pad_value, const char* mode, const mlx_stream s);
-
-int mlx_partition_axis(mlx_array* res, const mlx_array a, int kth, int axis, const mlx_stream s);
-
-int mlx_partition(mlx_array* res, const mlx_array a, int kth, const mlx_stream s);
-
-int mlx_power(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_prod_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s);
-
-int mlx_prod_axis(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s);
-
-int mlx_prod(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s);
-
-int mlx_put_along_axis(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array values, int axis, const mlx_stream s);
-
-int mlx_qqmm(mlx_array* res, const mlx_array x, const mlx_array w, const mlx_array w_scales , mlx_optional_int group_size, mlx_optional_int bits, const char* mode, const mlx_stream s);
-
-int mlx_quantize(mlx_vector_array* res, const mlx_array w, mlx_optional_int group_size, mlx_optional_int bits, const char* mode, const mlx_stream s);
-
-int mlx_quantized_matmul(mlx_array* res, const mlx_array x, const mlx_array w, const mlx_array scales, const mlx_array biases , bool transpose, mlx_optional_int group_size, mlx_optional_int bits, const char* mode, const mlx_stream s);
-
-int mlx_radians(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_real(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_reciprocal(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_remainder(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_repeat_axis(mlx_array* res, const mlx_array arr, int repeats, int axis, const mlx_stream s);
-
-int mlx_repeat(mlx_array* res, const mlx_array arr, int repeats, const mlx_stream s);
-
-int mlx_reshape(mlx_array* res, const mlx_array a, const int* shape, size_t shape_num, const mlx_stream s);
-
-int mlx_right_shift(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_roll_axis(mlx_array* res, const mlx_array a, const int* shift, size_t shift_num, int axis, const mlx_stream s);
-
-int mlx_roll_axes(mlx_array* res, const mlx_array a, const int* shift, size_t shift_num, const int* axes, size_t axes_num, const mlx_stream s);
-
-int mlx_roll(mlx_array* res, const mlx_array a, const int* shift, size_t shift_num, const mlx_stream s);
-
-int mlx_round(mlx_array* res, const mlx_array a, int decimals, const mlx_stream s);
-
-int mlx_rsqrt(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_scatter(mlx_array* res, const mlx_array a, const mlx_vector_array indices, const mlx_array updates, const int* axes, size_t axes_num, const mlx_stream s);
-
-int mlx_scatter_single(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array updates, int axis, const mlx_stream s);
-
-int mlx_scatter_add(mlx_array* res, const mlx_array a, const mlx_vector_array indices, const mlx_array updates, const int* axes, size_t axes_num, const mlx_stream s);
-
-int mlx_scatter_add_single(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array updates, int axis, const mlx_stream s);
-
-int mlx_scatter_add_axis(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array values, int axis, const mlx_stream s);
-
-int mlx_scatter_max(mlx_array* res, const mlx_array a, const mlx_vector_array indices, const mlx_array updates, const int* axes, size_t axes_num, const mlx_stream s);
-
-int mlx_scatter_max_single(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array updates, int axis, const mlx_stream s);
-
-int mlx_scatter_min(mlx_array* res, const mlx_array a, const mlx_vector_array indices, const mlx_array updates, const int* axes, size_t axes_num, const mlx_stream s);
-
-int mlx_scatter_min_single(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array updates, int axis, const mlx_stream s);
-
-int mlx_scatter_prod(mlx_array* res, const mlx_array a, const mlx_vector_array indices, const mlx_array updates, const int* axes, size_t axes_num, const mlx_stream s);
-
-int mlx_scatter_prod_single(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_array updates, int axis, const mlx_stream s);
-
-int mlx_segmented_mm(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_array segments, const mlx_stream s);
-
-int mlx_sigmoid(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_sign(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_sin(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_sinh(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_slice(mlx_array* res, const mlx_array a, const int* start, size_t start_num, const int* stop, size_t stop_num, const int* strides, size_t strides_num, const mlx_stream s);
-
-int mlx_slice_dynamic(mlx_array* res, const mlx_array a, const mlx_array start, const int* axes, size_t axes_num, const int* slice_size, size_t slice_size_num, const mlx_stream s);
-
-int mlx_slice_update(mlx_array* res, const mlx_array src, const mlx_array update, const int* start, size_t start_num, const int* stop, size_t stop_num, const int* strides, size_t strides_num, const mlx_stream s);
-
-int mlx_slice_update_dynamic(mlx_array* res, const mlx_array src, const mlx_array update, const mlx_array start, const int* axes, size_t axes_num, const mlx_stream s);
-
-int mlx_softmax_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool precise, const mlx_stream s);
-
-int mlx_softmax_axis(mlx_array* res, const mlx_array a, int axis, bool precise, const mlx_stream s);
-
-int mlx_softmax(mlx_array* res, const mlx_array a, bool precise, const mlx_stream s);
-
-int mlx_sort_axis(mlx_array* res, const mlx_array a, int axis, const mlx_stream s);
-
-int mlx_sort(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_split(mlx_vector_array* res, const mlx_array a, int num_splits, int axis, const mlx_stream s);
-
-int mlx_split_sections(mlx_vector_array* res, const mlx_array a, const int* indices, size_t indices_num, int axis, const mlx_stream s);
-
-int mlx_sqrt(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_square(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_squeeze_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, const mlx_stream s);
-
-int mlx_squeeze_axis(mlx_array* res, const mlx_array a, int axis, const mlx_stream s);
-
-int mlx_squeeze(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_stack_axis(mlx_array* res, const mlx_vector_array arrays, int axis, const mlx_stream s);
-
-int mlx_stack(mlx_array* res, const mlx_vector_array arrays, const mlx_stream s);
-
-int mlx_std_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, int ddof, const mlx_stream s);
-
-int mlx_std_axis(mlx_array* res, const mlx_array a, int axis, bool keepdims, int ddof, const mlx_stream s);
-
-int mlx_std(mlx_array* res, const mlx_array a, bool keepdims, int ddof, const mlx_stream s);
-
-int mlx_stop_gradient(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_subtract(mlx_array* res, const mlx_array a, const mlx_array b, const mlx_stream s);
-
-int mlx_sum_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s);
-
-int mlx_sum_axis(mlx_array* res, const mlx_array a, int axis, bool keepdims, const mlx_stream s);
-
-int mlx_sum(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s);
-
-int mlx_swapaxes(mlx_array* res, const mlx_array a, int axis1, int axis2, const mlx_stream s);
-
-int mlx_take_axis(mlx_array* res, const mlx_array a, const mlx_array indices, int axis, const mlx_stream s);
-
-int mlx_take(mlx_array* res, const mlx_array a, const mlx_array indices, const mlx_stream s);
-
-int mlx_take_along_axis(mlx_array* res, const mlx_array a, const mlx_array indices, int axis, const mlx_stream s);
-
-int mlx_tan(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_tanh(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_tensordot(mlx_array* res, const mlx_array a, const mlx_array b, const int* axes_a, size_t axes_a_num, const int* axes_b, size_t axes_b_num, const mlx_stream s);
-
-int mlx_tensordot_axis(mlx_array* res, const mlx_array a, const mlx_array b, int axis, const mlx_stream s);
-
-int mlx_tile(mlx_array* res, const mlx_array arr, const int* reps, size_t reps_num, const mlx_stream s);
-
-int mlx_to_fp8(mlx_array* res, const mlx_array x, const mlx_stream s);
-
-int mlx_topk_axis(mlx_array* res, const mlx_array a, int k, int axis, const mlx_stream s);
-
-int mlx_topk(mlx_array* res, const mlx_array a, int k, const mlx_stream s);
-
-int mlx_trace(mlx_array* res, const mlx_array a, int offset, int axis1, int axis2, mlx_dtype dtype, const mlx_stream s);
-
-int mlx_transpose_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, const mlx_stream s);
-
-int mlx_transpose(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_tri(mlx_array* res, int n, int m, int k, mlx_dtype type, const mlx_stream s);
-
-int mlx_tril(mlx_array* res, const mlx_array x, int k, const mlx_stream s);
-
-int mlx_triu(mlx_array* res, const mlx_array x, int k, const mlx_stream s);
-
-int mlx_unflatten(mlx_array* res, const mlx_array a, int axis, const int* shape, size_t shape_num, const mlx_stream s);
-
-int mlx_var_axes(mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, int ddof, const mlx_stream s);
-
-int mlx_var_axis(mlx_array* res, const mlx_array a, int axis, bool keepdims, int ddof, const mlx_stream s);
-
-int mlx_var(mlx_array* res, const mlx_array a, bool keepdims, int ddof, const mlx_stream s);
-
-int mlx_view(mlx_array* res, const mlx_array a, mlx_dtype dtype, const mlx_stream s);
-
-int mlx_where(mlx_array* res, const mlx_array condition, const mlx_array x, const mlx_array y, const mlx_stream s);
-
-int mlx_zeros(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_stream s);
-
-int mlx_zeros_like(mlx_array* res, const mlx_array a, const mlx_stream s);
-
-int mlx_random_bernoulli(mlx_array* res, const mlx_array p, const int* shape, size_t shape_num, const mlx_array key , const mlx_stream s);
-
-int mlx_random_bits(mlx_array* res, const int* shape, size_t shape_num, int width, const mlx_array key , const mlx_stream s);
-
-int mlx_random_categorical_shape(mlx_array* res, const mlx_array logits, int axis, const int* shape, size_t shape_num, const mlx_array key , const mlx_stream s);
-
-int mlx_random_categorical_num_samples(mlx_array* res, const mlx_array logits_, int axis, int num_samples, const mlx_array key , const mlx_stream s);
-
-int mlx_random_categorical(mlx_array* res, const mlx_array logits, int axis, const mlx_array key , const mlx_stream s);
-
-int mlx_random_gumbel(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_array key , const mlx_stream s);
-
-int mlx_random_key(mlx_array* res, uint64_t seed);
-
-int mlx_random_laplace(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, float loc, float scale, const mlx_array key , const mlx_stream s);
-
-int mlx_random_multivariate_normal(mlx_array* res, const mlx_array mean, const mlx_array cov, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_array key , const mlx_stream s);
-
-int mlx_random_normal_broadcast(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_array loc , const mlx_array scale , const mlx_array key , const mlx_stream s);
-
-int mlx_random_normal(mlx_array* res, const int* shape, size_t shape_num, mlx_dtype dtype, float loc, float scale, const mlx_array key , const mlx_stream s);
-
-int mlx_random_permutation(mlx_array* res, const mlx_array x, int axis, const mlx_array key , const mlx_stream s);
-
-int mlx_random_permutation_arange(mlx_array* res, int x, const mlx_array key , const mlx_stream s);
-
-int mlx_random_randint(mlx_array* res, const mlx_array low, const mlx_array high, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_array key , const mlx_stream s);
-
-int mlx_random_seed(uint64_t seed);
-
-int mlx_random_split_num(mlx_array* res, const mlx_array key, int num, const mlx_stream s);
-
-int mlx_random_split(mlx_array* res_0, mlx_array* res_1, const mlx_array key, const mlx_stream s);
-
-int mlx_random_truncated_normal(mlx_array* res, const mlx_array lower, const mlx_array upper, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_array key , const mlx_stream s);
-
-int mlx_random_uniform(mlx_array* res, const mlx_array low, const mlx_array high, const int* shape, size_t shape_num, mlx_dtype dtype, const mlx_array key , const mlx_stream s);
-
-mlx_stream mlx_stream_new(void);
-
-mlx_stream mlx_stream_new_device(mlx_device dev);
-
-int mlx_stream_set(mlx_stream* stream, const mlx_stream src);
-
-int mlx_stream_free(mlx_stream stream);
-
-int mlx_stream_tostring(mlx_string* str, mlx_stream stream);
-
-bool mlx_stream_equal(mlx_stream lhs, mlx_stream rhs);
-
-int mlx_stream_get_device(mlx_device* dev, mlx_stream stream);
-
-int mlx_stream_get_index(int* index, mlx_stream stream);
-
-int mlx_synchronize(mlx_stream stream);
-
-int mlx_get_default_stream(mlx_stream* stream, mlx_device dev);
-
-int mlx_set_default_stream(mlx_stream stream);
-
-mlx_stream mlx_default_cpu_stream_new(void);
-
-mlx_stream mlx_default_gpu_stream_new(void);
-
-mlx_string mlx_string_new(void);
-
-mlx_string mlx_string_new_data(const char* str);
-
-int mlx_string_set(mlx_string* str, const mlx_string src);
-
-const char* mlx_string_data(mlx_string str);
-
-int mlx_string_free(mlx_string str);
-
-int mlx_async_eval(const mlx_vector_array outputs);
-
-int mlx_checkpoint(mlx_closure* res, const mlx_closure fun);
-
-int mlx_custom_function(mlx_closure* res, const mlx_closure fun, const mlx_closure_custom fun_vjp , const mlx_closure_custom_jvp fun_jvp , const mlx_closure_custom_vmap fun_vmap);
-
-int mlx_custom_vjp(mlx_closure* res, const mlx_closure fun, const mlx_closure_custom fun_vjp);
-
-int mlx_eval(const mlx_vector_array outputs);
-
-int mlx_jvp(mlx_vector_array* res_0, mlx_vector_array* res_1, const mlx_closure fun, const mlx_vector_array primals, const mlx_vector_array tangents);
-
-int mlx_value_and_grad(mlx_closure_value_and_grad* res, const mlx_closure fun, const int* argnums, size_t argnums_num);
-
-int mlx_vjp(mlx_vector_array* res_0, mlx_vector_array* res_1, const mlx_closure fun, const mlx_vector_array primals, const mlx_vector_array cotangents);
-
-int mlx_detail_vmap_replace(mlx_vector_array* res, const mlx_vector_array inputs, const mlx_vector_array s_inputs, const mlx_vector_array s_outputs, const int* in_axes, size_t in_axes_num, const int* out_axes, size_t out_axes_num);
-
-int mlx_detail_vmap_trace(mlx_vector_array* res_0, mlx_vector_array* res_1, const mlx_closure fun, const mlx_vector_array inputs, const int* in_axes, size_t in_axes_num);
-
-mlx_vector_array mlx_vector_array_new(void);
-
-int mlx_vector_array_set(mlx_vector_array* vec, const mlx_vector_array src);
-
-int mlx_vector_array_free(mlx_vector_array vec);
-
-mlx_vector_array mlx_vector_array_new_data(const mlx_array* data, size_t size);
-
-mlx_vector_array mlx_vector_array_new_value(const mlx_array val);
-
-int mlx_vector_array_set_data(mlx_vector_array* vec, const mlx_array* data, size_t size);
-
-int mlx_vector_array_set_value(mlx_vector_array* vec, const mlx_array val);
-
-int mlx_vector_array_append_data(mlx_vector_array vec, const mlx_array* data, size_t size);
-
-int mlx_vector_array_append_value(mlx_vector_array vec, const mlx_array val);
-
-size_t mlx_vector_array_size(mlx_vector_array vec);
-
-int mlx_vector_array_get(mlx_array* res, const mlx_vector_array vec, size_t idx);
-
-mlx_vector_vector_array mlx_vector_vector_array_new(void);
-
-int mlx_vector_vector_array_set(mlx_vector_vector_array* vec, const mlx_vector_vector_array src);
-
-int mlx_vector_vector_array_free(mlx_vector_vector_array vec);
-
-mlx_vector_vector_array mlx_vector_vector_array_new_data(const mlx_vector_array* data, size_t size);
-
-mlx_vector_vector_array mlx_vector_vector_array_new_value(const mlx_vector_array val);
-
-int mlx_vector_vector_array_set_data(mlx_vector_vector_array* vec, const mlx_vector_array* data, size_t size);
-
-int mlx_vector_vector_array_set_value(mlx_vector_vector_array* vec, const mlx_vector_array val);
-
-int mlx_vector_vector_array_append_data(mlx_vector_vector_array vec, const mlx_vector_array* data, size_t size);
-
-int mlx_vector_vector_array_append_value(mlx_vector_vector_array vec, const mlx_vector_array val);
-
-size_t mlx_vector_vector_array_size(mlx_vector_vector_array vec);
-
-int mlx_vector_vector_array_get(mlx_vector_array* res, const mlx_vector_vector_array vec, size_t idx);
-
-mlx_vector_int mlx_vector_int_new(void);
-
-int mlx_vector_int_set(mlx_vector_int* vec, const mlx_vector_int src);
-
-int mlx_vector_int_free(mlx_vector_int vec);
-
-mlx_vector_int mlx_vector_int_new_data(int* data, size_t size);
-
-mlx_vector_int mlx_vector_int_new_value(int val);
-
-int mlx_vector_int_set_data(mlx_vector_int* vec, int* data, size_t size);
-
-int mlx_vector_int_set_value(mlx_vector_int* vec, int val);
-
-int mlx_vector_int_append_data(mlx_vector_int vec, int* data, size_t size);
-
-int mlx_vector_int_append_value(mlx_vector_int vec, int val);
-
-size_t mlx_vector_int_size(mlx_vector_int vec);
-
-int mlx_vector_int_get(int* res, const mlx_vector_int vec, size_t idx);
-
-mlx_vector_string mlx_vector_string_new(void);
-
-int mlx_vector_string_set(mlx_vector_string* vec, const mlx_vector_string src);
-
-int mlx_vector_string_free(mlx_vector_string vec);
-
-mlx_vector_string mlx_vector_string_new_data(const char** data, size_t size);
-
-mlx_vector_string mlx_vector_string_new_value(const char* val);
-
-int mlx_vector_string_set_data(mlx_vector_string* vec, const char** data, size_t size);
-
-int mlx_vector_string_set_value(mlx_vector_string* vec, const char* val);
-
-int mlx_vector_string_append_data(mlx_vector_string vec, const char** data, size_t size);
-
-int mlx_vector_string_append_value(mlx_vector_string vec, const char* val);
-
-size_t mlx_vector_string_size(mlx_vector_string vec);
-
-int mlx_vector_string_get(char** res, const mlx_vector_string vec, size_t idx);
-
-int mlx_version(mlx_string* str_);
-
-#endif // MLX_WRAPPERS_H

x/imagegen/mlx/mlx_dynamic.c

@@ -1,144 +0,0 @@
-// mlx_dynamic.c - Dynamic loading wrapper for MLX-C library
-// This file provides runtime dynamic loading of libmlxc instead of link-time binding
-
-#include "mlx_dynamic.h"
-#include <stdio.h>
-#include <stdlib.h>
-#include <string.h>
-
-#ifdef _WIN32
-#include <windows.h>
-typedef HMODULE lib_handle_t;
-#define LOAD_LIB(path) LoadLibraryA(path)
-#define GET_SYMBOL(handle, name) GetProcAddress(handle, name)
-#define CLOSE_LIB(handle) FreeLibrary(handle)
-#define LIB_ERROR() "LoadLibrary failed"
-#else
-#include <dlfcn.h>
-typedef void* lib_handle_t;
-#define LOAD_LIB(path) dlopen(path, RTLD_LAZY | RTLD_GLOBAL)
-#define GET_SYMBOL(handle, name) dlsym(handle, name)
-#define CLOSE_LIB(handle) dlclose(handle)
-#define LIB_ERROR() dlerror()
-#ifdef __APPLE__
-#include <mach-o/dyld.h>
-#include <libgen.h>
-#endif
-#endif
-
-static lib_handle_t mlx_handle = NULL;
-static int mlx_initialized = 0;
-static char mlx_error_buffer[512] = {0};
-
-#ifdef __APPLE__
-// Get path to library in same directory as executable
-static char* get_exe_relative_path(const char* libname) {
-    static char path[1024];
-    uint32_t size = sizeof(path);
-    if (_NSGetExecutablePath(path, &size) != 0) {
-        return NULL;
-    }
-    // Get directory of executable
-    char* dir = dirname(path);
-    static char fullpath[1024];
-    snprintf(fullpath, sizeof(fullpath), "%s/%s", dir, libname);
-    return fullpath;
-}
-#endif
-
-// Try to load library from a specific path
-static int try_load_lib(const char* path) {
-    if (!path) return 0;
-    mlx_handle = LOAD_LIB(path);
-    return mlx_handle != NULL;
-}
-
-// Initialize MLX dynamic library
-// Returns 0 on success, -1 on failure
-// On failure, call mlx_dynamic_error() to get error message
-int mlx_dynamic_init(void) {
-    if (mlx_initialized) {
-        return 0;  // Already initialized
-    }
-
-    const char* lib_path = NULL;
-    const char* tried_paths[8] = {0};
-    int num_tried = 0;
-
-#ifdef _WIN32
-    // Windows: try same directory as executable
-    lib_path = "libmlxc.dll";
-    tried_paths[num_tried++] = lib_path;
-    if (try_load_lib(lib_path)) goto success;
-#elif defined(__APPLE__)
-    // macOS: try executable directory first
-    lib_path = get_exe_relative_path("libmlxc.dylib");
-    if (lib_path) {
-        tried_paths[num_tried++] = lib_path;
-        if (try_load_lib(lib_path)) goto success;
-    }
-    // Try build directory (for tests run from repo root)
-    lib_path = "./build/lib/ollama/libmlxc.dylib";
-    tried_paths[num_tried++] = lib_path;
-    if (try_load_lib(lib_path)) goto success;
-    // Fallback to system paths
-    lib_path = "libmlxc.dylib";
-    tried_paths[num_tried++] = lib_path;
-    if (try_load_lib(lib_path)) goto success;
-#else
-    // Linux: try build directory first (for tests)
-    lib_path = "./build/lib/ollama/libmlxc.so";
-    tried_paths[num_tried++] = lib_path;
-    if (try_load_lib(lib_path)) goto success;
-    // Fallback to system paths
-    lib_path = "libmlxc.so";
-    tried_paths[num_tried++] = lib_path;
-    if (try_load_lib(lib_path)) goto success;
-#endif
-
-    // Failed to load library - build error message with all tried paths
-    {
-        const char* err = LIB_ERROR();
-        int offset = snprintf(mlx_error_buffer, sizeof(mlx_error_buffer),
-                     "MLX: Failed to load libmlxc library. Tried: ");
-        for (int i = 0; i < num_tried && offset < (int)sizeof(mlx_error_buffer) - 50; i++) {
-            offset += snprintf(mlx_error_buffer + offset, sizeof(mlx_error_buffer) - offset,
-                             "%s%s", i > 0 ? ", " : "", tried_paths[i]);
-        }
-        if (err) {
-            snprintf(mlx_error_buffer + offset, sizeof(mlx_error_buffer) - offset,
-                    ". Last error: %s", err);
-        }
-    }
-    return -1;
-
-success:
-    mlx_initialized = 1;
-    snprintf(mlx_error_buffer, sizeof(mlx_error_buffer),
-             "MLX: Successfully loaded %s", lib_path ? lib_path : "library");
-    return 0;
-}
-
-// Get the last error message
-const char* mlx_dynamic_error(void) {
-    return mlx_error_buffer;
-}
-
-// Check if MLX is initialized
-int mlx_dynamic_is_initialized(void) {
-    return mlx_initialized;
-}
-
-// Get the library handle (for use by generated wrappers)
-void* mlx_get_handle(void) {
-    return mlx_handle;
-}
-
-// Cleanup (optional, called at program exit)
-void mlx_dynamic_cleanup(void) {
-    if (mlx_handle != NULL) {
-        CLOSE_LIB(mlx_handle);
-        mlx_handle = NULL;
-        mlx_initialized = 0;
-    }
-}

x/imagegen/mlx/mlx_dynamic.h

@@ -1,29 +0,0 @@
-// mlx_dynamic.h - Dynamic loading interface for MLX-C library
-#ifndef MLX_DYNAMIC_H
-#define MLX_DYNAMIC_H
-
-#ifdef __cplusplus
-extern "C" {
-#endif
-
-// Initialize the MLX dynamic library
-// Returns 0 on success, -1 on failure
-int mlx_dynamic_init(void);
-
-// Get the last error message from dynamic loading
-const char* mlx_dynamic_error(void);
-
-// Check if MLX is initialized
-int mlx_dynamic_is_initialized(void);
-
-// Get the library handle (for use by generated wrappers)
-void* mlx_get_handle(void);
-
-// Cleanup resources (optional, for clean shutdown)
-void mlx_dynamic_cleanup(void);
-
-#ifdef __cplusplus
-}
-#endif
-
-#endif // MLX_DYNAMIC_H

x/imagegen/mlx/mlx_test.go

@@ -4,30 +4,9 @@ package mlx
 
 import (
 	"fmt"
-	"os"
-	"path/filepath"
-	"runtime"
 	"testing"
 )
 
-// TestMain initializes MLX before running tests.
-// If MLX libraries are not available, tests are skipped.
-func TestMain(m *testing.M) {
-	// Change to repo root so ./build/lib/ollama/ path works
-	_, thisFile, _, _ := runtime.Caller(0)
-	repoRoot := filepath.Join(filepath.Dir(thisFile), "..", "..", "..")
-	if err := os.Chdir(repoRoot); err != nil {
-		fmt.Printf("Failed to change to repo root: %v\n", err)
-		os.Exit(1)
-	}
-
-	if err := InitMLX(); err != nil {
-		fmt.Printf("Skipping MLX tests: %v\n", err)
-		os.Exit(0)
-	}
-	os.Exit(m.Run())
-}
-
 // TestBasicCleanup verifies non-kept arrays are freed and kept arrays survive.
 func TestBasicCleanup(t *testing.T) {
 	weight := NewArrayFloat32([]float32{1, 2, 3, 4}, []int32{2, 2})

x/imagegen/models/flux2/flux2.go

@@ -1,539 +0,0 @@
-//go:build mlx
-
-// Package flux2 implements the FLUX.2 Klein diffusion transformer model.
-// Klein is a 4B parameter distilled model that supports sub-second inference.
-package flux2
-
-import (
-	"context"
-	"encoding/json"
-	"fmt"
-	"image"
-	"math"
-	"time"
-
-	"github.com/ollama/ollama/x/imagegen"
-	"github.com/ollama/ollama/x/imagegen/mlx"
-	"github.com/ollama/ollama/x/imagegen/models/qwen3"
-	"github.com/ollama/ollama/x/imagegen/tokenizer"
-	"golang.org/x/image/draw"
-)
-
-// GenerateConfig holds all options for image generation.
-type GenerateConfig struct {
-	Prompt        string
-	Width         int32                 // Image width (default: 1024)
-	Height        int32                 // Image height (default: 1024)
-	Steps         int                   // Denoising steps (default: 4 for Klein)
-	GuidanceScale float32               // Guidance scale (default: 1.0, Klein doesn't need CFG)
-	Seed          int64                 // Random seed
-	Progress      func(step, totalSteps int) // Optional progress callback
-	CapturePath   string                // GPU capture path (debug)
-	InputImages   []image.Image         // Reference images for image conditioning (already loaded)
-}
-
-// Model represents a FLUX.2 Klein model.
-type Model struct {
-	ModelName       string
-	Tokenizer       *tokenizer.Tokenizer
-	TextEncoder     *qwen3.TextEncoder
-	Transformer     *Flux2Transformer2DModel
-	VAE             *AutoencoderKLFlux2
-	SchedulerConfig *SchedulerConfig
-}
-
-// TextEncoderLayerIndices are the layers from which to extract text embeddings.
-// Diffusers uses hidden_states[9, 18, 27]. In Python, hidden_states[0] is the embedding
-// output before any layers, so hidden_states[9] = after layer 8 (0-indexed).
-// Go's ForwardWithLayerOutputs captures after layer i runs, so we use [8, 17, 26].
-var TextEncoderLayerIndices = []int{8, 17, 26}
-
-// Load loads the FLUX.2 Klein model from ollama blob storage.
-func (m *Model) Load(modelName string) error {
-	fmt.Printf("Loading FLUX.2 Klein model from manifest: %s...\n", modelName)
-	start := time.Now()
-
-	if mlx.GPUIsAvailable() {
-		mlx.SetDefaultDeviceGPU()
-		mlx.EnableCompile()
-	}
-
-	m.ModelName = modelName
-
-	// Load manifest
-	manifest, err := imagegen.LoadManifest(modelName)
-	if err != nil {
-		return fmt.Errorf("load manifest: %w", err)
-	}
-
-	// Load tokenizer
-	fmt.Print("  Loading tokenizer... ")
-	tokData, err := manifest.ReadConfig("tokenizer/tokenizer.json")
-	if err != nil {
-		return fmt.Errorf("tokenizer: %w", err)
-	}
-
-	tokConfig := &tokenizer.TokenizerConfig{}
-	if data, err := manifest.ReadConfig("tokenizer/tokenizer_config.json"); err == nil {
-		tokConfig.TokenizerConfigJSON = data
-	}
-	if data, err := manifest.ReadConfig("tokenizer/generation_config.json"); err == nil {
-		tokConfig.GenerationConfigJSON = data
-	}
-	if data, err := manifest.ReadConfig("tokenizer/special_tokens_map.json"); err == nil {
-		tokConfig.SpecialTokensMapJSON = data
-	}
-
-	tok, err := tokenizer.LoadFromBytesWithConfig(tokData, tokConfig)
-	if err != nil {
-		return fmt.Errorf("tokenizer: %w", err)
-	}
-	m.Tokenizer = tok
-	fmt.Println("✓")
-
-	// Load text encoder
-	m.TextEncoder = &qwen3.TextEncoder{}
-	if err := m.TextEncoder.Load(manifest, "text_encoder/config.json"); err != nil {
-		return fmt.Errorf("text encoder: %w", err)
-	}
-
-	// Load transformer
-	m.Transformer = &Flux2Transformer2DModel{}
-	if err := m.Transformer.Load(manifest); err != nil {
-		return fmt.Errorf("transformer: %w", err)
-	}
-
-	// Load VAE
-	m.VAE = &AutoencoderKLFlux2{}
-	if err := m.VAE.Load(manifest); err != nil {
-		return fmt.Errorf("VAE: %w", err)
-	}
-
-	// Evaluate all weights in a single batch (reduces GPU sync overhead)
-	fmt.Print("  Evaluating weights... ")
-	allWeights := mlx.Collect(m.TextEncoder)
-	allWeights = append(allWeights, mlx.Collect(m.Transformer)...)
-	allWeights = append(allWeights, mlx.Collect(m.VAE)...)
-	mlx.Eval(allWeights...)
-	fmt.Println("✓")
-
-	// Load scheduler config
-	m.SchedulerConfig = DefaultSchedulerConfig()
-	if schedData, err := manifest.ReadConfig("scheduler/scheduler_config.json"); err == nil {
-		if err := json.Unmarshal(schedData, m.SchedulerConfig); err != nil {
-			fmt.Printf("  Warning: failed to parse scheduler config: %v\n", err)
-		}
-	}
-
-	mem := mlx.MetalGetActiveMemory()
-	fmt.Printf("  Loaded in %.2fs (%.1f GB VRAM)\n", time.Since(start).Seconds(), float64(mem)/(1024*1024*1024))
-
-	return nil
-}
-
-// Generate creates an image from a prompt.
-func (m *Model) Generate(prompt string, width, height int32, steps int, seed int64) (*mlx.Array, error) {
-	return m.GenerateFromConfig(context.Background(), &GenerateConfig{
-		Prompt: prompt,
-		Width:  width,
-		Height: height,
-		Steps:  steps,
-		Seed:   seed,
-	})
-}
-
-// GenerateWithProgress creates an image with progress callback.
-func (m *Model) GenerateWithProgress(prompt string, width, height int32, steps int, seed int64, progress func(step, totalSteps int)) (*mlx.Array, error) {
-	return m.GenerateFromConfig(context.Background(), &GenerateConfig{
-		Prompt:   prompt,
-		Width:    width,
-		Height:   height,
-		Steps:    steps,
-		Seed:     seed,
-		Progress: progress,
-	})
-}
-
-// GenerateFromConfig generates an image using the unified config struct.
-func (m *Model) GenerateFromConfig(ctx context.Context, cfg *GenerateConfig) (*mlx.Array, error) {
-	start := time.Now()
-	result, err := m.generate(ctx, cfg)
-	if err != nil {
-		return nil, err
-	}
-	fmt.Printf("Generated in %.2fs (%d steps)\n", time.Since(start).Seconds(), cfg.Steps)
-	return result, nil
-}
-
-// GenerateImage implements runner.ImageModel interface.
-func (m *Model) GenerateImage(ctx context.Context, prompt string, width, height int32, steps int, seed int64, progress func(step, total int)) (*mlx.Array, error) {
-	return m.GenerateFromConfig(ctx, &GenerateConfig{
-		Prompt:   prompt,
-		Width:    width,
-		Height:   height,
-		Steps:    steps,
-		Seed:     seed,
-		Progress: progress,
-	})
-}
-
-// MaxOutputPixels is the maximum output resolution (4 megapixels, ~2048x2048)
-const MaxOutputPixels = 2048 * 2048
-
-// MaxRefPixels is the maximum resolution for reference images (smaller to reduce attention memory)
-const MaxRefPixels = 728 * 728
-
-// generate is the internal denoising pipeline.
-func (m *Model) generate(ctx context.Context, cfg *GenerateConfig) (*mlx.Array, error) {
-	// Enable MLX compilation for fused kernels
-	mlx.EnableCompile()
-
-	// Apply defaults
-	if cfg.Steps <= 0 {
-		cfg.Steps = 4 // Klein default: 4 steps for distilled model
-	}
-	if cfg.GuidanceScale <= 0 {
-		cfg.GuidanceScale = 1.0 // Klein doesn't need guidance
-	}
-
-	// Determine output dimensions
-	if len(cfg.InputImages) > 0 {
-		// With input images, compute missing dimension from aspect ratio
-		// Images are already EXIF-rotated by the caller
-		bounds := cfg.InputImages[0].Bounds()
-		imgW, imgH := bounds.Dx(), bounds.Dy()
-		aspectRatio := float64(imgH) / float64(imgW)
-		if cfg.Width > 0 && cfg.Height <= 0 {
-			// Width specified, compute height
-			cfg.Height = int32(math.Round(float64(cfg.Width)*aspectRatio/16) * 16)
-		} else if cfg.Height > 0 && cfg.Width <= 0 {
-			// Height specified, compute width
-			cfg.Width = int32(math.Round(float64(cfg.Height)/aspectRatio/16) * 16)
-		} else if cfg.Width <= 0 && cfg.Height <= 0 {
-			// Neither specified, use input dimensions
-			cfg.Width = int32(imgW)
-			cfg.Height = int32(imgH)
-		}
-	}
-	if cfg.Width <= 0 {
-		cfg.Width = 1024
-	}
-	if cfg.Height <= 0 {
-		cfg.Height = 1024
-	}
-
-	// Cap to max pixels, preserve aspect ratio, round to multiple of 16
-	pixels := int(cfg.Width) * int(cfg.Height)
-	if pixels > MaxOutputPixels {
-		scale := math.Sqrt(float64(MaxOutputPixels) / float64(pixels))
-		cfg.Width = int32(math.Round(float64(cfg.Width) * scale / 16) * 16)
-		cfg.Height = int32(math.Round(float64(cfg.Height) * scale / 16) * 16)
-	}
-	cfg.Height = int32((cfg.Height + 8) / 16 * 16) // round to nearest 16
-	cfg.Width = int32((cfg.Width + 8) / 16 * 16)
-	fmt.Printf("  Output: %dx%d\n", cfg.Width, cfg.Height)
-
-	tcfg := m.Transformer.TransformerConfig
-	patchSize := m.VAE.Config.PatchSize
-
-	// Latent dimensions: image / 8 (VAE downscale) / patch_size
-	latentH := cfg.Height / 8
-	latentW := cfg.Width / 8
-	patchH := latentH / patchSize[0]
-	patchW := latentW / patchSize[1]
-	imgSeqLen := patchH * patchW
-
-	// Text encoding with multi-layer extraction (no padding, use true sequence length)
-	fmt.Print("  Encoding prompt... ")
-	promptEmbeds, textLen := m.TextEncoder.EncodePromptWithLayers(m.Tokenizer, cfg.Prompt, 512, TextEncoderLayerIndices, false)
-	fmt.Println("✓")
-
-	// Encode reference images if provided
-	var refTokens *ImageCondTokens
-	var refHeights, refWidths []int32
-	if len(cfg.InputImages) > 0 {
-		fmt.Printf("  Encoding %d reference image(s):\n", len(cfg.InputImages))
-
-		var err error
-		refTokens, err = m.EncodeImageRefs(cfg.InputImages)
-		if err != nil {
-			return nil, fmt.Errorf("encode reference images: %w", err)
-		}
-
-		// Extract heights/widths for RoPE computation (same limits as EncodeImageRefs)
-		limitPixels := MaxRefPixels
-		if len(cfg.InputImages) > 1 {
-			limitPixels = MaxRefPixels / 2
-		}
-		for _, img := range cfg.InputImages {
-			_, w, h := PrepareImage(img, limitPixels)
-			refHeights = append(refHeights, int32(h/16))
-			refWidths = append(refWidths, int32(w/16))
-		}
-	}
-
-	// Scheduler
-	scheduler := NewFlowMatchScheduler(m.SchedulerConfig)
-	scheduler.SetTimestepsWithMu(cfg.Steps, CalculateShift(imgSeqLen, cfg.Steps))
-
-	// Init latents in packed form [B, C*4, H/2, W/2] like diffusers
-	// diffusers creates noise in [B, 128, 64, 64] and packs to [B, 4096, 128]
-	latentChannels := m.VAE.Config.LatentChannels
-	packedChannels := latentChannels * 4 // 32 * 4 = 128
-	latents := scheduler.InitNoise([]int32{1, packedChannels, patchH, patchW}, cfg.Seed)
-
-	// Pack latents (transpose): [B, C, H, W] -> [B, H*W, C]
-	// This matches diffusers' _pack_latents
-	patches := packLatents(latents)
-	noiseSeqLen := patches.Shape()[1]
-
-	// RoPE cache - includes reference images if present
-	rope := PrepareRoPECache(textLen, patchH, patchW, tcfg.AxesDimsRoPE, tcfg.RopeTheta, refHeights, refWidths, ImageRefScale)
-
-	// Cleanup setup arrays when done
-	defer func() {
-		rope.Cos.Free()
-		rope.Sin.Free()
-		promptEmbeds.Free()
-		if refTokens != nil {
-			refTokens.Tokens.Free()
-		}
-	}()
-
-	// Pre-compute all timesteps before the loop to avoid per-step tensor creation
-	timesteps := make([]*mlx.Array, cfg.Steps)
-	for i := 0; i < cfg.Steps; i++ {
-		tCurr := scheduler.Timesteps[i] / float32(m.SchedulerConfig.NumTrainTimesteps)
-		timesteps[i] = mlx.ToBFloat16(mlx.NewArray([]float32{tCurr}, []int32{1}))
-	}
-
-	// Evaluate setup arrays
-	fmt.Print("  Evaluating setup... ")
-	setupStart := time.Now()
-	toEval := []*mlx.Array{promptEmbeds, patches, rope.Cos, rope.Sin}
-	toEval = append(toEval, timesteps...)
-	if refTokens != nil {
-		toEval = append(toEval, refTokens.Tokens)
-	}
-	mlx.Eval(toEval...)
-	mlx.MetalResetPeakMemory() // Reset peak to measure generation separately
-	fmt.Printf("✓ (%.2fs, %.1f GB)\n", time.Since(setupStart).Seconds(),
-		float64(mlx.MetalGetActiveMemory())/(1024*1024*1024))
-
-	if cfg.Progress != nil {
-		cfg.Progress(0, cfg.Steps)
-	}
-
-	loopStart := time.Now()
-	stepStart := time.Now()
-
-	// Denoising loop
-	for i := 0; i < cfg.Steps; i++ {
-		// Check for cancellation
-		if ctx != nil {
-			select {
-			case <-ctx.Done():
-				return nil, ctx.Err()
-			default:
-			}
-		}
-
-		// GPU capture on step 2 if requested
-		if cfg.CapturePath != "" && i == 1 {
-			mlx.MetalStartCapture(cfg.CapturePath)
-		}
-
-		timestep := timesteps[i]
-
-		// Prepare input - concatenate noise patches with reference tokens if present
-		imgInput := patches
-		if refTokens != nil {
-			imgInput = mlx.Concatenate([]*mlx.Array{patches, refTokens.Tokens}, 1)
-		}
-
-		// Transformer forward pass
-		output := m.Transformer.Forward(imgInput, promptEmbeds, timestep, rope)
-
-		// If we concatenated reference tokens, slice to only get noise portion
-		if refTokens != nil {
-			output = mlx.Slice(output, []int32{0, 0, 0}, []int32{1, noiseSeqLen, output.Shape()[2]})
-		}
-
-		// Scheduler step (keep reference to old patches for the computation graph)
-		newPatches := scheduler.Step(output, patches, i)
-
-		if cfg.CapturePath != "" && i == 1 {
-			mlx.MetalStopCapture()
-		}
-
-		mlx.Eval(newPatches)
-		patches = newPatches
-
-		elapsed := time.Since(stepStart).Seconds()
-		peakGB := float64(mlx.MetalGetPeakMemory()) / (1024 * 1024 * 1024)
-		if i == 0 {
-			fmt.Printf("    step %d: %.2fs (JIT warmup), peak %.1f GB\n", i+1, elapsed, peakGB)
-		} else {
-			fmt.Printf("    step %d: %.2fs, peak %.1f GB\n", i+1, elapsed, peakGB)
-		}
-		stepStart = time.Now()
-		if cfg.Progress != nil {
-			cfg.Progress(i+1, cfg.Steps)
-		}
-	}
-
-	loopTime := time.Since(loopStart).Seconds()
-	peakMem := float64(mlx.MetalGetPeakMemory()) / (1024 * 1024 * 1024)
-	fmt.Printf("  Denoised %d steps in %.2fs (%.2fs/step), peak %.1f GB\n",
-		cfg.Steps, loopTime, loopTime/float64(cfg.Steps), peakMem)
-
-	// Free timesteps now that denoising is done
-	for _, ts := range timesteps {
-		ts.Free()
-	}
-
-	// VAE decode with tiling for larger images
-	fmt.Print("  Decoding VAE... ")
-	vaeStart := time.Now()
-	// Enable tiling for images > 512x512 (latent > 64x64)
-	// VAE attention is O(n²) on latent pixels, tiling reduces memory significantly
-	if patchH*2 > 64 || patchW*2 > 64 {
-		m.VAE.Tiling = DefaultTilingConfig()
-	}
-	decoded := m.VAE.Decode(patches, patchH, patchW)
-	mlx.Eval(decoded)
-
-	// Free patches now that decode is done
-	patches.Free()
-
-	fmt.Printf("✓ (%.2fs, peak %.1f GB)\n", time.Since(vaeStart).Seconds(),
-		float64(mlx.MetalGetPeakMemory())/(1024*1024*1024))
-
-	return decoded, nil
-}
-
-// packLatents converts [B, C, H, W] to [B, H*W, C] (matches diffusers _pack_latents)
-func packLatents(x *mlx.Array) *mlx.Array {
-	shape := x.Shape()
-	B := shape[0]
-	C := shape[1]
-	H := shape[2]
-	W := shape[3]
-	// [B, C, H, W] -> [B, C, H*W] -> [B, H*W, C]
-	x = mlx.Reshape(x, B, C, H*W)
-	return mlx.Transpose(x, 0, 2, 1)
-}
-
-// LoadPersistent loads the model and keeps it in memory for repeated use.
-func LoadPersistent(modelName string) (*Model, error) {
-	m := &Model{}
-	if err := m.Load(modelName); err != nil {
-		return nil, err
-	}
-	return m, nil
-}
-
-// ImageRefScale is the time coordinate offset between reference images (matches diffusers scale=10)
-const ImageRefScale = 10
-
-// PrepareImage resizes and crops an image to be a multiple of 16, with optional pixel limit.
-// Returns the processed image and its dimensions.
-func PrepareImage(img image.Image, limitPixels int) (image.Image, int, int) {
-	bounds := img.Bounds()
-	w, h := bounds.Dx(), bounds.Dy()
-
-	// Cap pixels if needed (like diffusers cap_pixels)
-	if limitPixels > 0 && w*h > limitPixels {
-		scale := math.Sqrt(float64(limitPixels) / float64(w*h))
-		w = int(float64(w) * scale)
-		h = int(float64(h) * scale)
-	}
-
-	// Round down to multiple of 16
-	w = (w / 16) * 16
-	h = (h / 16) * 16
-
-	if w < 16 {
-		w = 16
-	}
-	if h < 16 {
-		h = 16
-	}
-
-	// Resize using high-quality bicubic interpolation (matches diffusers' default lanczos)
-	resized := image.NewRGBA(image.Rect(0, 0, w, h))
-	draw.CatmullRom.Scale(resized, resized.Bounds(), img, img.Bounds(), draw.Over, nil)
-
-	return resized, w, h
-}
-
-// ImageToTensor converts an image to a tensor in [-1, 1] range with shape [1, C, H, W].
-func ImageToTensor(img image.Image) *mlx.Array {
-	bounds := img.Bounds()
-	w, h := bounds.Dx(), bounds.Dy()
-
-	// Convert to float32 array in NCHW format [1, 3, H, W] with values in [-1, 1]
-	data := make([]float32, 3*h*w)
-
-	for y := 0; y < h; y++ {
-		for x := 0; x < w; x++ {
-			r, g, b, _ := img.At(x+bounds.Min.X, y+bounds.Min.Y).RGBA()
-			// RGBA returns 16-bit values, convert to [-1, 1]
-			data[0*h*w+y*w+x] = float32(r>>8)/127.5 - 1.0
-			data[1*h*w+y*w+x] = float32(g>>8)/127.5 - 1.0
-			data[2*h*w+y*w+x] = float32(b>>8)/127.5 - 1.0
-		}
-	}
-
-	arr := mlx.NewArrayFloat32(data, []int32{1, 3, int32(h), int32(w)})
-	return arr
-}
-
-// ImageCondTokens holds encoded reference image tokens.
-type ImageCondTokens struct {
-	Tokens *mlx.Array // [1, total_tokens, C] - concatenated reference tokens
-}
-
-// EncodeImageRefs encodes reference images using the VAE.
-func (m *Model) EncodeImageRefs(images []image.Image) (*ImageCondTokens, error) {
-	if len(images) == 0 {
-		return nil, nil
-	}
-
-	// Limit reference images to reduce attention memory
-	limitPixels := MaxRefPixels
-	if len(images) > 1 {
-		limitPixels = MaxRefPixels / 2
-	}
-
-	var allTokens []*mlx.Array
-
-	for _, img := range images {
-		// Prepare image (resize, crop to multiple of 16)
-		prepared, prepW, prepH := PrepareImage(img, limitPixels)
-		fmt.Printf("    Encoding %dx%d image... ", prepW, prepH)
-
-		// Convert to tensor [-1, 1]
-		tensor := ImageToTensor(prepared)
-
-		// Encode with VAE - returns [1, L, 128]
-		encoded := m.VAE.EncodeImage(tensor)
-		squeezed := mlx.Squeeze(encoded, 0) // [L, C]
-
-		// Defer eval - will be done with other setup arrays
-		allTokens = append(allTokens, squeezed)
-		fmt.Println("✓")
-	}
-
-	// For single image, just add batch dimension directly
-	// For multiple images, concatenate first
-	var tokens *mlx.Array
-	if len(allTokens) == 1 {
-		tokens = mlx.ExpandDims(allTokens[0], 0) // [1, L, C]
-	} else {
-		tokens = mlx.Concatenate(allTokens, 0) // [total_L, C]
-		tokens = mlx.ExpandDims(tokens, 0)     // [1, total_L, C]
-	}
-
-	return &ImageCondTokens{Tokens: tokens}, nil
-}

x/imagegen/models/flux2/rope.go

@@ -1,224 +0,0 @@
-//go:build mlx
-
-package flux2
-
-import (
-	"math"
-
-	"github.com/ollama/ollama/x/imagegen/mlx"
-)
-
-// RoPEConfig holds 4D RoPE configuration for Flux2
-type RoPEConfig struct {
-	Theta    int32   // 2000 for Klein
-	AxesDims []int32 // [32, 32, 32, 32] - dimensions for T, H, W, L axes
-}
-
-// RoPECache holds precomputed RoPE cos/sin values
-type RoPECache struct {
-	Cos      *mlx.Array // [1, TotalSeqLen, 1, head_dim/2]
-	Sin      *mlx.Array // [1, TotalSeqLen, 1, head_dim/2]
-	TextLen  int32      // Length of text sequence
-	ImageLen int32      // Length of image sequence
-}
-
-// PrepareTextIDs creates position IDs for text tokens.
-// Text tokens use: T=0, H=0, W=0, L=0..seqLen-1
-// Returns: [seqLen, 4]
-func PrepareTextIDs(seqLen int32) *mlx.Array {
-	ids := make([]float32, seqLen*4)
-	for i := int32(0); i < seqLen; i++ {
-		idx := i * 4
-		ids[idx+0] = 0             // T = 0
-		ids[idx+1] = 0             // H = 0
-		ids[idx+2] = 0             // W = 0
-		ids[idx+3] = float32(i)    // L = sequence position
-	}
-	return mlx.NewArray(ids, []int32{seqLen, 4})
-}
-
-// PrepareLatentIDs creates position IDs for image latent tokens.
-// Latent tokens use: T=0, H=0..height-1, W=0..width-1, L=0
-// The latents are in row-major order (H then W).
-// Returns: [height*width, 4]
-func PrepareLatentIDs(height, width int32) *mlx.Array {
-	seqLen := height * width
-	ids := make([]float32, seqLen*4)
-	idx := 0
-	for h := int32(0); h < height; h++ {
-		for w := int32(0); w < width; w++ {
-			ids[idx*4+0] = 0           // T = 0
-			ids[idx*4+1] = float32(h)  // H = row
-			ids[idx*4+2] = float32(w)  // W = column
-			ids[idx*4+3] = 0           // L = 0
-			idx++
-		}
-	}
-	return mlx.NewArray(ids, []int32{seqLen, 4})
-}
-
-// PrepareImageIDs creates position IDs for reference image tokens (used in editing).
-// Reference images use: T=scale*(i+1), H=0..h-1, W=0..w-1, L=0
-// where i is the image index (0, 1, 2, ...) and scale separates images in T dimension.
-// Returns: [total_tokens, 4]
-func PrepareImageIDs(imageHeights, imageWidths []int32, scale int32) *mlx.Array {
-	// Calculate total tokens
-	totalTokens := int32(0)
-	for i := range imageHeights {
-		totalTokens += imageHeights[i] * imageWidths[i]
-	}
-
-	ids := make([]float32, totalTokens*4)
-	idx := int32(0)
-	for imgIdx, h := range imageHeights {
-		w := imageWidths[imgIdx]
-		tValue := float32(scale * int32(imgIdx+1))
-		for hi := int32(0); hi < h; hi++ {
-			for wi := int32(0); wi < w; wi++ {
-				ids[idx*4+0] = tValue       // T = scale * (imgIdx + 1)
-				ids[idx*4+1] = float32(hi)  // H = row
-				ids[idx*4+2] = float32(wi)  // W = column
-				ids[idx*4+3] = 0            // L = 0
-				idx++
-			}
-		}
-	}
-	return mlx.NewArray(ids, []int32{totalTokens, 4})
-}
-
-// ComputeRoPE computes cos and sin for 4D rotary position embeddings.
-// ids: [L, 4] with (T, H, W, L) coordinates
-// axesDims: [32, 32, 32, 32] - each axis has this many dimensions (total = head_dim = 128)
-// theta: base frequency (2000 for Klein)
-// Returns: cos, sin each [1, L, 1, head_dim] with repeat_interleave applied
-func ComputeRoPE(ids *mlx.Array, axesDims []int32, theta int32) (*mlx.Array, *mlx.Array) {
-	shape := ids.Shape()
-	seqLen := shape[0]
-
-	// Compute total head dim (sum of all axes dims)
-	headDim := int32(0)
-	for _, d := range axesDims {
-		headDim += d
-	}
-
-	// Extract each coordinate dimension
-	// ids[:, 0] = T, ids[:, 1] = H, ids[:, 2] = W, ids[:, 3] = L
-	posT := mlx.Slice(ids, []int32{0, 0}, []int32{seqLen, 1}) // [L, 1]
-	posH := mlx.Slice(ids, []int32{0, 1}, []int32{seqLen, 2}) // [L, 1]
-	posW := mlx.Slice(ids, []int32{0, 2}, []int32{seqLen, 3}) // [L, 1]
-	posL := mlx.Slice(ids, []int32{0, 3}, []int32{seqLen, 4}) // [L, 1]
-
-	// Compute frequencies for each axis
-	logTheta := float32(math.Log(float64(theta)))
-	cosArrs := make([]*mlx.Array, 4)
-	sinArrs := make([]*mlx.Array, 4)
-	positions := []*mlx.Array{posT, posH, posW, posL}
-
-	for i, axisDim := range axesDims {
-		half := axisDim / 2
-
-		// Create frequency array for this axis: theta^(-2j/dim) for j=0..half-1
-		// This matches diffusers: 1.0 / (theta ** (torch.arange(0, dim, 2) / dim))
-		freqs := make([]float32, half)
-		for j := int32(0); j < half; j++ {
-			freqs[j] = float32(math.Exp(float64(-logTheta * float32(2*j) / float32(axisDim))))
-		}
-		freqArr := mlx.NewArray(freqs, []int32{1, half})
-
-		// Compute pos * freq -> [L, half]
-		posExpanded := positions[i] // [L, 1]
-		args := mlx.Mul(posExpanded, freqArr) // [L, half]
-
-		// Compute cos and sin for this axis
-		cosAxis := mlx.Cos(args) // [L, half]
-		sinAxis := mlx.Sin(args) // [L, half]
-
-		// repeat_interleave(2): [c0, c1, ...] -> [c0, c0, c1, c1, ...]
-		// Reshape [L, half] -> [L, half, 1], tile to [L, half, 2], reshape to [L, axisDim]
-		cosAxis = mlx.ExpandDims(cosAxis, 2)                        // [L, half, 1]
-		cosAxis = mlx.Tile(cosAxis, []int32{1, 1, 2})               // [L, half, 2]
-		cosAxis = mlx.Reshape(cosAxis, seqLen, axisDim)             // [L, axisDim]
-
-		sinAxis = mlx.ExpandDims(sinAxis, 2)
-		sinAxis = mlx.Tile(sinAxis, []int32{1, 1, 2})
-		sinAxis = mlx.Reshape(sinAxis, seqLen, axisDim)
-
-		cosArrs[i] = cosAxis
-		sinArrs[i] = sinAxis
-	}
-
-	// Concatenate all axes: [L, headDim]
-	cos := mlx.Concatenate(cosArrs, 1)
-	sin := mlx.Concatenate(sinArrs, 1)
-
-	// Reshape to [1, L, 1, headDim] for broadcasting with attention
-	cos = mlx.Reshape(cos, 1, seqLen, 1, headDim)
-	sin = mlx.Reshape(sin, 1, seqLen, 1, headDim)
-
-	return cos, sin
-}
-
-// ApplyRoPE4D applies 4D rotary position embeddings to queries and keys.
-// x: [B, L, nheads, head_dim]
-// cos, sin: [1, L, 1, head_dim] (with repeat_interleave applied)
-// Returns: x with RoPE applied
-// Matches diffusers apply_rotary_emb with use_real=True, use_real_unbind_dim=-1
-func ApplyRoPE4D(x *mlx.Array, cos, sin *mlx.Array) *mlx.Array {
-	shape := x.Shape()
-	B := shape[0]
-	L := shape[1]
-	nheads := shape[2]
-	headDim := shape[3]
-	half := headDim / 2
-
-	// Reshape x to [B, L, nheads, half, 2] and split into real/imag
-	xReshaped := mlx.Reshape(x, B, L, nheads, half, 2)
-
-	// Extract real (index 0) and imag (index 1) parts
-	xReal := mlx.Slice(xReshaped, []int32{0, 0, 0, 0, 0}, []int32{B, L, nheads, half, 1})
-	xImag := mlx.Slice(xReshaped, []int32{0, 0, 0, 0, 1}, []int32{B, L, nheads, half, 2})
-	xReal = mlx.Squeeze(xReal, 4) // [B, L, nheads, half]
-	xImag = mlx.Squeeze(xImag, 4) // [B, L, nheads, half]
-
-	// x_rotated = stack([-x_imag, x_real], dim=-1).flatten(-2)
-	// This creates [-x_imag[0], x_real[0], -x_imag[1], x_real[1], ...]
-	negXImag := mlx.Neg(xImag)
-	negXImag = mlx.ExpandDims(negXImag, 4) // [B, L, nheads, half, 1]
-	xReal = mlx.ExpandDims(xReal, 4)       // [B, L, nheads, half, 1]
-	xRotated := mlx.Concatenate([]*mlx.Array{negXImag, xReal}, 4) // [B, L, nheads, half, 2]
-	xRotated = mlx.Reshape(xRotated, B, L, nheads, headDim)       // [B, L, nheads, headDim]
-
-	// out = x * cos + x_rotated * sin
-	return mlx.Add(mlx.Mul(x, cos), mlx.Mul(xRotated, sin))
-}
-
-// PrepareRoPECache creates RoPE cache for text + noise, optionally with reference images.
-// textLen: number of text tokens
-// noiseH, noiseW: dimensions of the noise latent in patch tokens
-// axesDims: [32, 32, 32, 32]
-// theta: 2000
-// refHeights, refWidths: optional reference image dimensions (pass nil/empty for no images)
-// scale: time coordinate offset between reference images (e.g., 10)
-func PrepareRoPECache(textLen, noiseH, noiseW int32, axesDims []int32, theta int32, refHeights, refWidths []int32, scale int32) *RoPECache {
-	textIDs := PrepareTextIDs(textLen)
-	noiseIDs := PrepareLatentIDs(noiseH, noiseW)
-
-	var allIDs *mlx.Array
-	imageLen := noiseH * noiseW
-
-	if len(refHeights) > 0 {
-		refIDs := PrepareImageIDs(refHeights, refWidths, scale)
-		allIDs = mlx.Concatenate([]*mlx.Array{textIDs, noiseIDs, refIDs}, 0)
-		for i := range refHeights {
-			imageLen += refHeights[i] * refWidths[i]
-		}
-	} else {
-		allIDs = mlx.Concatenate([]*mlx.Array{textIDs, noiseIDs}, 0)
-	}
-
-	cos, sin := ComputeRoPE(allIDs, axesDims, theta)
-	cos = mlx.ToBFloat16(cos)
-	sin = mlx.ToBFloat16(sin)
-
-	return &RoPECache{Cos: cos, Sin: sin, TextLen: textLen, ImageLen: imageLen}
-}

x/imagegen/models/flux2/scheduler.go

@@ -1,149 +0,0 @@
-//go:build mlx
-
-package flux2
-
-import (
-	"math"
-
-	"github.com/ollama/ollama/x/imagegen/mlx"
-)
-
-// SchedulerConfig holds Flow-Match scheduler configuration
-type SchedulerConfig struct {
-	NumTrainTimesteps  int32   `json:"num_train_timesteps"`  // 1000
-	Shift              float32 `json:"shift"`                // 3.0 for Klein
-	UseDynamicShifting bool    `json:"use_dynamic_shifting"` // true
-	TimeShiftType      string  `json:"time_shift_type"`      // "exponential" or "linear"
-}
-
-// DefaultSchedulerConfig returns default config for Klein
-func DefaultSchedulerConfig() *SchedulerConfig {
-	return &SchedulerConfig{
-		NumTrainTimesteps:  1000,
-		Shift:              3.0, // Klein uses 3.0
-		UseDynamicShifting: true,
-		TimeShiftType:      "exponential",
-	}
-}
-
-// FlowMatchScheduler implements the Flow-Match Euler discrete scheduler
-type FlowMatchScheduler struct {
-	Config    *SchedulerConfig
-	Timesteps []float32 // Discretized timesteps (t from 1 to 0)
-	Sigmas    []float32 // Noise levels at each timestep
-	NumSteps  int       // Number of inference steps
-}
-
-// NewFlowMatchScheduler creates a new scheduler
-func NewFlowMatchScheduler(cfg *SchedulerConfig) *FlowMatchScheduler {
-	return &FlowMatchScheduler{
-		Config: cfg,
-	}
-}
-
-// SetTimesteps sets up the scheduler for the given number of inference steps
-func (s *FlowMatchScheduler) SetTimesteps(numSteps int) {
-	s.SetTimestepsWithMu(numSteps, 0)
-}
-
-// SetTimestepsWithMu sets up scheduler matching diffusers set_timesteps(sigmas=..., mu=...)
-func (s *FlowMatchScheduler) SetTimestepsWithMu(numSteps int, mu float32) {
-	s.NumSteps = numSteps
-
-	// diffusers: sigmas = linspace(1, 1/num_steps, num_steps)
-	// Then applies time shift, appends 0.0 at end
-	s.Sigmas = make([]float32, numSteps+1)
-
-	for i := 0; i < numSteps; i++ {
-		// linspace(1, 1/num_steps, num_steps)
-		var sigma float32
-		if numSteps == 1 {
-			sigma = 1.0
-		} else {
-			sigma = 1.0 - float32(i)/float32(numSteps-1)*(1.0-1.0/float32(numSteps))
-		}
-
-		// Apply time shift if using dynamic shifting
-		if s.Config.UseDynamicShifting && mu != 0 {
-			sigma = s.timeShift(mu, sigma)
-		} else {
-			// If not dynamic shifting, apply fixed shift scaling like diffusers
-			shift := s.Config.Shift
-			sigma = shift * sigma / (1 + (shift-1)*sigma)
-		}
-		s.Sigmas[i] = sigma
-	}
-	// Append terminal zero
-	s.Sigmas[numSteps] = 0.0
-
-	// Timesteps scaled to training range (matches diffusers: timesteps = sigmas * num_train_timesteps)
-	s.Timesteps = make([]float32, numSteps+1)
-	for i, v := range s.Sigmas {
-		s.Timesteps[i] = v * float32(s.Config.NumTrainTimesteps)
-	}
-}
-
-// timeShift applies the dynamic time shift
-func (s *FlowMatchScheduler) timeShift(mu float32, t float32) float32 {
-	if t <= 0 {
-		return 0
-	}
-	if s.Config.TimeShiftType == "linear" {
-		return mu / (mu + (1.0/t-1.0))
-	}
-	// Default: exponential
-	expMu := float32(math.Exp(float64(mu)))
-	return expMu / (expMu + (1.0/t - 1.0))
-}
-
-// Step performs one denoising step
-func (s *FlowMatchScheduler) Step(modelOutput, sample *mlx.Array, timestepIdx int) *mlx.Array {
-	sigma := s.Sigmas[timestepIdx]
-	sigmaNext := s.Sigmas[timestepIdx+1]
-
-	// Euler step: x_{t-dt} = x_t + (sigma_next - sigma) * v_t
-	dt := sigmaNext - sigma
-
-	// Upcast to float32 for precision (matches diffusers)
-	sampleF32 := mlx.AsType(sample, mlx.DtypeFloat32)
-	outputF32 := mlx.AsType(modelOutput, mlx.DtypeFloat32)
-
-	scaledOutput := mlx.MulScalar(outputF32, dt)
-	result := mlx.Add(sampleF32, scaledOutput)
-
-	// Cast back to bfloat16
-	return mlx.ToBFloat16(result)
-}
-
-// GetTimestep returns the timestep value at the given index
-func (s *FlowMatchScheduler) GetTimestep(idx int) float32 {
-	if idx < len(s.Timesteps) {
-		return s.Timesteps[idx]
-	}
-	return 0.0
-}
-
-// InitNoise creates initial noise for sampling
-func (s *FlowMatchScheduler) InitNoise(shape []int32, seed int64) *mlx.Array {
-	return mlx.RandomNormalWithDtype(shape, uint64(seed), mlx.DtypeBFloat16)
-}
-
-// CalculateShift computes the mu shift value for dynamic scheduling
-// Matches diffusers compute_empirical_mu function
-func CalculateShift(imgSeqLen int32, numSteps int) float32 {
-	a1, b1 := float32(8.73809524e-05), float32(1.89833333)
-	a2, b2 := float32(0.00016927), float32(0.45666666)
-
-	seqLen := float32(imgSeqLen)
-
-	if imgSeqLen > 4300 {
-		return a2*seqLen + b2
-	}
-
-	m200 := a2*seqLen + b2
-	m10 := a1*seqLen + b1
-
-	a := (m200 - m10) / 190.0
-	b := m200 - 200.0*a
-	return a*float32(numSteps) + b
-}

x/imagegen/models/flux2/transformer.go

@@ -1,562 +0,0 @@
-//go:build mlx
-
-package flux2
-
-import (
-	"fmt"
-	"math"
-
-	"github.com/ollama/ollama/x/imagegen"
-	"github.com/ollama/ollama/x/imagegen/mlx"
-	"github.com/ollama/ollama/x/imagegen/nn"
-	"github.com/ollama/ollama/x/imagegen/safetensors"
-)
-
-// TransformerConfig holds Flux2 transformer configuration
-type TransformerConfig struct {
-	AttentionHeadDim         int32   `json:"attention_head_dim"`          // 128
-	AxesDimsRoPE             []int32 `json:"axes_dims_rope"`              // [32, 32, 32, 32]
-	Eps                      float32 `json:"eps"`                         // 1e-6
-	GuidanceEmbeds           bool    `json:"guidance_embeds"`             // false for Klein
-	InChannels               int32   `json:"in_channels"`                 // 128
-	JointAttentionDim        int32   `json:"joint_attention_dim"`         // 7680
-	MLPRatio                 float32 `json:"mlp_ratio"`                   // 3.0
-	NumAttentionHeads        int32   `json:"num_attention_heads"`         // 24
-	NumLayers                int32   `json:"num_layers"`                  // 5
-	NumSingleLayers          int32   `json:"num_single_layers"`           // 20
-	PatchSize                int32   `json:"patch_size"`                  // 1
-	RopeTheta                int32   `json:"rope_theta"`                  // 2000
-	TimestepGuidanceChannels int32   `json:"timestep_guidance_channels"`  // 256
-}
-
-// Computed dimensions
-func (c *TransformerConfig) InnerDim() int32 {
-	return c.NumAttentionHeads * c.AttentionHeadDim // 24 * 128 = 3072
-}
-
-func (c *TransformerConfig) MLPHiddenDim() int32 {
-	return int32(float32(c.InnerDim()) * c.MLPRatio) // 3072 * 3.0 = 9216
-}
-
-// TimestepEmbedder creates timestep embeddings
-// Weight names: time_guidance_embed.timestep_embedder.linear_1.weight, linear_2.weight
-type TimestepEmbedder struct {
-	Linear1  nn.LinearLayer `weight:"linear_1"`
-	Linear2  nn.LinearLayer `weight:"linear_2"`
-	EmbedDim int32          // 256
-}
-
-// Forward creates sinusoidal embeddings and projects them
-func (t *TimestepEmbedder) Forward(timesteps *mlx.Array) *mlx.Array {
-	half := t.EmbedDim / 2
-	freqs := make([]float32, half)
-	for i := int32(0); i < half; i++ {
-		freqs[i] = float32(math.Exp(-math.Log(10000.0) * float64(i) / float64(half)))
-	}
-	freqsArr := mlx.NewArray(freqs, []int32{1, half})
-
-	// timesteps: [B] -> [B, 1]
-	tExpanded := mlx.ExpandDims(timesteps, 1)
-	// args: [B, half]
-	args := mlx.Mul(tExpanded, freqsArr)
-
-	// [cos(args), sin(args)] -> [B, embed_dim]
-	sinEmbed := mlx.Concatenate([]*mlx.Array{mlx.Cos(args), mlx.Sin(args)}, 1)
-
-	// MLP: linear_1 -> silu -> linear_2
-	h := t.Linear1.Forward(sinEmbed)
-	h = mlx.SiLU(h)
-	return t.Linear2.Forward(h)
-}
-
-// TimeGuidanceEmbed wraps the timestep embedder
-// Weight names: time_guidance_embed.timestep_embedder.*
-type TimeGuidanceEmbed struct {
-	TimestepEmbedder *TimestepEmbedder `weight:"timestep_embedder"`
-}
-
-// Forward computes timestep embeddings
-func (t *TimeGuidanceEmbed) Forward(timesteps *mlx.Array) *mlx.Array {
-	return t.TimestepEmbedder.Forward(timesteps)
-}
-
-// Modulation computes adaptive modulation parameters
-// Weight names: double_stream_modulation_img.linear.weight, etc.
-type Modulation struct {
-	Linear nn.LinearLayer `weight:"linear"`
-}
-
-// Forward computes modulation parameters
-func (m *Modulation) Forward(temb *mlx.Array) *mlx.Array {
-	h := mlx.SiLU(temb)
-	return m.Linear.Forward(h)
-}
-
-// TransformerBlockAttn implements dual-stream attention
-// Weight names: transformer_blocks.N.attn.*
-type TransformerBlockAttn struct {
-	// Image stream (separate Q, K, V projections)
-	ToQ nn.LinearLayer `weight:"to_q"`
-	ToK nn.LinearLayer `weight:"to_k"`
-	ToV nn.LinearLayer `weight:"to_v"`
-	// Note: to_out has .0 suffix in weights, handled specially
-	ToOut0 nn.LinearLayer `weight:"to_out.0"`
-
-	// Text stream (add_ projections)
-	AddQProj nn.LinearLayer `weight:"add_q_proj"`
-	AddKProj nn.LinearLayer `weight:"add_k_proj"`
-	AddVProj nn.LinearLayer `weight:"add_v_proj"`
-	ToAddOut nn.LinearLayer `weight:"to_add_out"`
-
-	// QK norms for image stream
-	NormQ *mlx.Array `weight:"norm_q.weight"`
-	NormK *mlx.Array `weight:"norm_k.weight"`
-
-	// QK norms for text stream (added)
-	NormAddedQ *mlx.Array `weight:"norm_added_q.weight"`
-	NormAddedK *mlx.Array `weight:"norm_added_k.weight"`
-}
-
-// FeedForward implements SwiGLU MLP
-// Weight names: transformer_blocks.N.ff.linear_in.weight, linear_out.weight
-type FeedForward struct {
-	LinearIn  nn.LinearLayer `weight:"linear_in"`
-	LinearOut nn.LinearLayer `weight:"linear_out"`
-}
-
-// Forward applies SwiGLU MLP
-func (ff *FeedForward) Forward(x *mlx.Array) *mlx.Array {
-	// LinearIn outputs 2x hidden dim for SwiGLU
-	h := ff.LinearIn.Forward(x)
-	shape := h.Shape()
-	half := shape[len(shape)-1] / 2
-
-	// Split into gate and up
-	gate := mlx.Slice(h, []int32{0, 0, 0}, []int32{shape[0], shape[1], half})
-	up := mlx.Slice(h, []int32{0, 0, half}, []int32{shape[0], shape[1], shape[2]})
-
-	// SwiGLU: silu(gate) * up
-	h = mlx.Mul(mlx.SiLU(gate), up)
-	return ff.LinearOut.Forward(h)
-}
-
-// TransformerBlock implements a dual-stream transformer block
-// Weight names: transformer_blocks.N.*
-type TransformerBlock struct {
-	Attn      *TransformerBlockAttn `weight:"attn"`
-	FF        *FeedForward          `weight:"ff"`
-	FFContext *FeedForward          `weight:"ff_context"`
-
-	// Config (set after loading)
-	NHeads  int32
-	HeadDim int32
-	Scale   float32
-}
-
-// Forward applies the dual-stream block
-// imgHidden: [B, imgLen, dim]
-// txtHidden: [B, txtLen, dim]
-// imgMod, txtMod: modulation params [B, 6*dim] each
-// cos, sin: RoPE values
-func (block *TransformerBlock) Forward(imgHidden, txtHidden *mlx.Array, imgMod, txtMod *mlx.Array, cos, sin *mlx.Array) (*mlx.Array, *mlx.Array) {
-	imgShape := imgHidden.Shape()
-	B := imgShape[0]
-	imgLen := imgShape[1]
-	dim := imgShape[2]
-	txtLen := txtHidden.Shape()[1]
-
-	// Parse modulation: 6 params each (shift1, scale1, gate1, shift2, scale2, gate2)
-	imgShift1, imgScale1, imgGate1 := parseModulation3(imgMod, dim, 0)
-	imgShift2, imgScale2, imgGate2 := parseModulation3(imgMod, dim, 3)
-	txtShift1, txtScale1, txtGate1 := parseModulation3(txtMod, dim, 0)
-	txtShift2, txtScale2, txtGate2 := parseModulation3(txtMod, dim, 3)
-
-	// === Attention branch ===
-	// Modulate inputs
-	imgNorm := modulateLayerNorm(imgHidden, imgShift1, imgScale1)
-	txtNorm := modulateLayerNorm(txtHidden, txtShift1, txtScale1)
-
-	// Compute Q, K, V for image stream (separate projections)
-	imgQ := block.Attn.ToQ.Forward(imgNorm)
-	imgK := block.Attn.ToK.Forward(imgNorm)
-	imgV := block.Attn.ToV.Forward(imgNorm)
-
-	// Compute Q, K, V for text stream (add_ projections)
-	txtQ := block.Attn.AddQProj.Forward(txtNorm)
-	txtK := block.Attn.AddKProj.Forward(txtNorm)
-	txtV := block.Attn.AddVProj.Forward(txtNorm)
-
-	// Reshape for attention: [B, L, dim] -> [B, L, nheads, headDim]
-	imgQ = mlx.Reshape(imgQ, B, imgLen, block.NHeads, block.HeadDim)
-	imgK = mlx.Reshape(imgK, B, imgLen, block.NHeads, block.HeadDim)
-	imgV = mlx.Reshape(imgV, B, imgLen, block.NHeads, block.HeadDim)
-	txtQ = mlx.Reshape(txtQ, B, txtLen, block.NHeads, block.HeadDim)
-	txtK = mlx.Reshape(txtK, B, txtLen, block.NHeads, block.HeadDim)
-	txtV = mlx.Reshape(txtV, B, txtLen, block.NHeads, block.HeadDim)
-
-	// Apply QK norm (RMSNorm with learned scale)
-	imgQ = applyQKNorm(imgQ, block.Attn.NormQ)
-	imgK = applyQKNorm(imgK, block.Attn.NormK)
-	txtQ = applyQKNorm(txtQ, block.Attn.NormAddedQ)
-	txtK = applyQKNorm(txtK, block.Attn.NormAddedK)
-
-	// Concatenate for joint attention: text first, then image
-	q := mlx.Concatenate([]*mlx.Array{txtQ, imgQ}, 1)
-	k := mlx.Concatenate([]*mlx.Array{txtK, imgK}, 1)
-	v := mlx.Concatenate([]*mlx.Array{txtV, imgV}, 1)
-
-	// Apply RoPE
-	q = ApplyRoPE4D(q, cos, sin)
-	k = ApplyRoPE4D(k, cos, sin)
-
-	// Transpose for SDPA: [B, nheads, L, headDim]
-	q = mlx.Transpose(q, 0, 2, 1, 3)
-	k = mlx.Transpose(k, 0, 2, 1, 3)
-	v = mlx.Transpose(v, 0, 2, 1, 3)
-
-	// Scaled dot-product attention
-	out := mlx.ScaledDotProductAttention(q, k, v, block.Scale, false)
-
-	// Transpose back: [B, L, nheads, headDim]
-	out = mlx.Transpose(out, 0, 2, 1, 3)
-
-	// Split back into txt and img
-	totalLen := txtLen + imgLen
-	txtOut := mlx.Slice(out, []int32{0, 0, 0, 0}, []int32{B, txtLen, block.NHeads, block.HeadDim})
-	imgOut := mlx.Slice(out, []int32{0, txtLen, 0, 0}, []int32{B, totalLen, block.NHeads, block.HeadDim})
-
-	// Reshape and project
-	txtOut = mlx.Reshape(txtOut, B, txtLen, dim)
-	imgOut = mlx.Reshape(imgOut, B, imgLen, dim)
-	txtOut = block.Attn.ToAddOut.Forward(txtOut)
-	imgOut = block.Attn.ToOut0.Forward(imgOut)
-
-	// Apply gates and residual
-	imgHidden = mlx.Add(imgHidden, mlx.Mul(imgGate1, imgOut))
-	txtHidden = mlx.Add(txtHidden, mlx.Mul(txtGate1, txtOut))
-
-	// === MLP branch ===
-	imgNorm = modulateLayerNorm(imgHidden, imgShift2, imgScale2)
-	txtNorm = modulateLayerNorm(txtHidden, txtShift2, txtScale2)
-
-	imgFFOut := block.FF.Forward(imgNorm)
-	txtFFOut := block.FFContext.Forward(txtNorm)
-
-	imgHidden = mlx.Add(imgHidden, mlx.Mul(imgGate2, imgFFOut))
-	txtHidden = mlx.Add(txtHidden, mlx.Mul(txtGate2, txtFFOut))
-
-	return imgHidden, txtHidden
-}
-
-// SingleTransformerBlockAttn implements attention for single-stream blocks
-// Weight names: single_transformer_blocks.N.attn.*
-type SingleTransformerBlockAttn struct {
-	ToQKVMlpProj nn.LinearLayer `weight:"to_qkv_mlp_proj"` // Fused QKV + MLP input
-	ToOut        nn.LinearLayer `weight:"to_out"`          // Fused attn_out + MLP out
-	NormQ        *mlx.Array     `weight:"norm_q.weight"`
-	NormK        *mlx.Array     `weight:"norm_k.weight"`
-}
-
-// SingleTransformerBlock implements a single-stream transformer block
-// Weight names: single_transformer_blocks.N.*
-type SingleTransformerBlock struct {
-	Attn *SingleTransformerBlockAttn `weight:"attn"`
-
-	// Config
-	NHeads    int32
-	HeadDim   int32
-	InnerDim  int32
-	MLPHidDim int32
-	Scale     float32
-}
-
-// Forward applies the single-stream block
-// x: [B, L, dim] concatenated text+image
-// mod: modulation [B, 3*dim]
-func (block *SingleTransformerBlock) Forward(x *mlx.Array, mod *mlx.Array, cos, sin *mlx.Array) *mlx.Array {
-	shape := x.Shape()
-	B := shape[0]
-	L := shape[1]
-	dim := shape[2]
-
-	// Parse modulation: (shift, scale, gate)
-	shift, scale, gate := parseModulation3(mod, dim, 0)
-
-	// Modulate input
-	h := modulateLayerNorm(x, shift, scale)
-
-	// Fused projection: QKV + MLP gate/up
-	// linear1 outputs: [q, k, v, mlp_gate, mlp_up] = [dim, dim, dim, mlpHid, mlpHid]
-	qkvMlp := block.Attn.ToQKVMlpProj.Forward(h)
-
-	// Split: first 3*dim is QKV, rest is MLP
-	qkvDim := 3 * block.InnerDim
-	qkv := mlx.Slice(qkvMlp, []int32{0, 0, 0}, []int32{B, L, qkvDim})
-	mlpIn := mlx.Slice(qkvMlp, []int32{0, 0, qkvDim}, []int32{B, L, qkvMlp.Shape()[2]})
-
-	// Split QKV
-	q, k, v := splitQKV(qkv, B, L, block.InnerDim)
-
-	// Reshape for attention
-	q = mlx.Reshape(q, B, L, block.NHeads, block.HeadDim)
-	k = mlx.Reshape(k, B, L, block.NHeads, block.HeadDim)
-	v = mlx.Reshape(v, B, L, block.NHeads, block.HeadDim)
-
-	// QK norm
-	q = applyQKNorm(q, block.Attn.NormQ)
-	k = applyQKNorm(k, block.Attn.NormK)
-
-	// Apply RoPE
-	q = ApplyRoPE4D(q, cos, sin)
-	k = ApplyRoPE4D(k, cos, sin)
-
-	// Transpose for SDPA
-	q = mlx.Transpose(q, 0, 2, 1, 3)
-	k = mlx.Transpose(k, 0, 2, 1, 3)
-	v = mlx.Transpose(v, 0, 2, 1, 3)
-
-	// SDPA
-	attnOut := mlx.ScaledDotProductAttention(q, k, v, block.Scale, false)
-
-	// Transpose back and reshape
-	attnOut = mlx.Transpose(attnOut, 0, 2, 1, 3)
-	attnOut = mlx.Reshape(attnOut, B, L, block.InnerDim)
-
-	// MLP: SwiGLU
-	mlpShape := mlpIn.Shape()
-	half := mlpShape[2] / 2
-	mlpGate := mlx.Slice(mlpIn, []int32{0, 0, 0}, []int32{B, L, half})
-	mlpUp := mlx.Slice(mlpIn, []int32{0, 0, half}, []int32{B, L, mlpShape[2]})
-	mlpOut := mlx.Mul(mlx.SiLU(mlpGate), mlpUp)
-
-	// Concatenate attention and MLP for fused output
-	combined := mlx.Concatenate([]*mlx.Array{attnOut, mlpOut}, 2)
-
-	// Output projection
-	out := block.Attn.ToOut.Forward(combined)
-
-	// Apply gate and residual
-	return mlx.Add(x, mlx.Mul(gate, out))
-}
-
-// NormOut implements the output normalization with modulation
-// Weight names: norm_out.linear.weight
-type NormOut struct {
-	Linear nn.LinearLayer `weight:"linear"`
-}
-
-// Forward computes final modulated output
-func (n *NormOut) Forward(x *mlx.Array, temb *mlx.Array) *mlx.Array {
-	shape := x.Shape()
-	B := shape[0]
-	dim := shape[2]
-
-	// Modulation: temb -> silu -> linear -> [shift, scale]
-	mod := mlx.SiLU(temb)
-	mod = n.Linear.Forward(mod)
-
-	// Split into scale and shift (diffusers order: scale first, shift second)
-	scale := mlx.Slice(mod, []int32{0, 0}, []int32{B, dim})
-	shift := mlx.Slice(mod, []int32{0, dim}, []int32{B, 2 * dim})
-	shift = mlx.ExpandDims(shift, 1)
-	scale = mlx.ExpandDims(scale, 1)
-
-	// Modulate with RMSNorm
-	return modulateLayerNorm(x, shift, scale)
-}
-
-// Flux2Transformer2DModel is the main Flux2 transformer
-// Weight names at top level: time_guidance_embed.*, double_stream_modulation_*.*, etc.
-type Flux2Transformer2DModel struct {
-	// Timestep embedding
-	TimeGuidanceEmbed *TimeGuidanceEmbed `weight:"time_guidance_embed"`
-
-	// Shared modulation
-	DoubleStreamModulationImg *Modulation `weight:"double_stream_modulation_img"`
-	DoubleStreamModulationTxt *Modulation `weight:"double_stream_modulation_txt"`
-	SingleStreamModulation    *Modulation `weight:"single_stream_modulation"`
-
-	// Embedders
-	XEmbedder       nn.LinearLayer `weight:"x_embedder"`
-	ContextEmbedder nn.LinearLayer `weight:"context_embedder"`
-
-	// Transformer blocks
-	TransformerBlocks       []*TransformerBlock       `weight:"transformer_blocks"`
-	SingleTransformerBlocks []*SingleTransformerBlock `weight:"single_transformer_blocks"`
-
-	// Output
-	NormOut *NormOut       `weight:"norm_out"`
-	ProjOut nn.LinearLayer `weight:"proj_out"`
-
-	*TransformerConfig
-}
-
-// Load loads the Flux2 transformer from ollama blob storage.
-func (m *Flux2Transformer2DModel) Load(manifest *imagegen.ModelManifest) error {
-	fmt.Print("  Loading transformer... ")
-
-	// Load config from blob
-	var cfg TransformerConfig
-	if err := manifest.ReadConfigJSON("transformer/config.json", &cfg); err != nil {
-		return fmt.Errorf("config: %w", err)
-	}
-	m.TransformerConfig = &cfg
-
-	// Initialize slices
-	m.TransformerBlocks = make([]*TransformerBlock, cfg.NumLayers)
-	m.SingleTransformerBlocks = make([]*SingleTransformerBlock, cfg.NumSingleLayers)
-
-	// Initialize TimeGuidanceEmbed with embed dim
-	m.TimeGuidanceEmbed = &TimeGuidanceEmbed{
-		TimestepEmbedder: &TimestepEmbedder{EmbedDim: cfg.TimestepGuidanceChannels},
-	}
-
-	// Load weights from tensor blobs
-	weights, err := imagegen.LoadWeightsFromManifest(manifest, "transformer")
-	if err != nil {
-		return fmt.Errorf("weights: %w", err)
-	}
-	if err := weights.Load(0); err != nil {
-		return fmt.Errorf("load weights: %w", err)
-	}
-	defer weights.ReleaseAll()
-
-	return m.loadWeights(weights)
-}
-
-// loadWeights loads weights from any WeightSource into the model
-func (m *Flux2Transformer2DModel) loadWeights(weights safetensors.WeightSource) error {
-	if err := safetensors.LoadModule(m, weights, ""); err != nil {
-		return fmt.Errorf("load module: %w", err)
-	}
-	m.initComputedFields()
-	fmt.Println("✓")
-	return nil
-}
-
-// initComputedFields initializes computed fields after loading weights
-func (m *Flux2Transformer2DModel) initComputedFields() {
-	cfg := m.TransformerConfig
-	innerDim := cfg.InnerDim()
-	scale := float32(1.0 / math.Sqrt(float64(cfg.AttentionHeadDim)))
-
-	// Initialize transformer blocks
-	for _, block := range m.TransformerBlocks {
-		block.NHeads = cfg.NumAttentionHeads
-		block.HeadDim = cfg.AttentionHeadDim
-		block.Scale = scale
-	}
-
-	// Initialize single transformer blocks
-	for _, block := range m.SingleTransformerBlocks {
-		block.NHeads = cfg.NumAttentionHeads
-		block.HeadDim = cfg.AttentionHeadDim
-		block.InnerDim = innerDim
-		block.MLPHidDim = cfg.MLPHiddenDim()
-		block.Scale = scale
-	}
-}
-
-// Forward runs the Flux2 transformer
-func (m *Flux2Transformer2DModel) Forward(patches, txtEmbeds *mlx.Array, timesteps *mlx.Array, rope *RoPECache) *mlx.Array {
-	patchShape := patches.Shape()
-	B := patchShape[0]
-	imgLen := patchShape[1]
-	txtLen := txtEmbeds.Shape()[1]
-
-	// Scale timestep to 0-1000 range (diffusers multiplies by 1000)
-	scaledTimesteps := mlx.MulScalar(timesteps, 1000.0)
-
-	// Compute timestep embedding
-	temb := m.TimeGuidanceEmbed.Forward(scaledTimesteps)
-
-	// Embed patches and text
-	imgHidden := m.XEmbedder.Forward(patches)
-	txtHidden := m.ContextEmbedder.Forward(txtEmbeds)
-
-	// Compute shared modulation
-	imgMod := m.DoubleStreamModulationImg.Forward(temb)
-	txtMod := m.DoubleStreamModulationTxt.Forward(temb)
-	singleMod := m.SingleStreamModulation.Forward(temb)
-
-	// Double (dual-stream) blocks
-	for _, block := range m.TransformerBlocks {
-		imgHidden, txtHidden = block.Forward(imgHidden, txtHidden, imgMod, txtMod, rope.Cos, rope.Sin)
-	}
-
-	// Concatenate for single-stream: text first, then image
-	hidden := mlx.Concatenate([]*mlx.Array{txtHidden, imgHidden}, 1)
-
-	// Single-stream blocks
-	for _, block := range m.SingleTransformerBlocks {
-		hidden = block.Forward(hidden, singleMod, rope.Cos, rope.Sin)
-	}
-
-	// Extract image portion
-	totalLen := txtLen + imgLen
-	imgOut := mlx.Slice(hidden, []int32{0, txtLen, 0}, []int32{B, totalLen, hidden.Shape()[2]})
-
-	// Final norm and projection
-	imgOut = m.NormOut.Forward(imgOut, temb)
-	return m.ProjOut.Forward(imgOut)
-}
-
-// Note: QK normalization uses mlx.RMSNorm (the fast version) directly
-// See applyQKNorm function below
-
-// compiledSwiGLU fuses: silu(gate) * up
-// Called 30x per step (10 in dual-stream + 20 in single-stream blocks)
-var compiledSwiGLU *mlx.CompiledFunc
-
-func getCompiledSwiGLU() *mlx.CompiledFunc {
-	if compiledSwiGLU == nil {
-		compiledSwiGLU = mlx.CompileShapeless(func(inputs []*mlx.Array) []*mlx.Array {
-			gate, up := inputs[0], inputs[1]
-			return []*mlx.Array{mlx.Mul(mlx.SiLU(gate), up)}
-		}, true)
-	}
-	return compiledSwiGLU
-}
-
-// Helper functions
-
-// parseModulation3 extracts 3 modulation params (shift, scale, gate) starting at offset
-func parseModulation3(mod *mlx.Array, dim int32, offset int32) (*mlx.Array, *mlx.Array, *mlx.Array) {
-	B := mod.Shape()[0]
-	start := offset * dim
-	shift := mlx.Slice(mod, []int32{0, start}, []int32{B, start + dim})
-	scale := mlx.Slice(mod, []int32{0, start + dim}, []int32{B, start + 2*dim})
-	gate := mlx.Slice(mod, []int32{0, start + 2*dim}, []int32{B, start + 3*dim})
-
-	// Expand for broadcasting [B, dim] -> [B, 1, dim]
-	shift = mlx.ExpandDims(shift, 1)
-	scale = mlx.ExpandDims(scale, 1)
-	gate = mlx.ExpandDims(gate, 1)
-
-	return shift, scale, gate
-}
-
-// modulateLayerNorm applies LayerNorm then shift/scale modulation
-// Diffusers uses LayerNorm(elementwise_affine=False) which centers the data
-func modulateLayerNorm(x *mlx.Array, shift, scale *mlx.Array) *mlx.Array {
-	// Fast LayerNorm without learnable params
-	x = mlx.LayerNorm(x, 1e-6)
-
-	// Modulate: x * (1 + scale) + shift
-	x = mlx.Mul(x, mlx.AddScalar(scale, 1.0))
-	return mlx.Add(x, shift)
-}
-
-// splitQKV splits a fused QKV tensor into Q, K, V
-func splitQKV(qkv *mlx.Array, B, L, dim int32) (*mlx.Array, *mlx.Array, *mlx.Array) {
-	q := mlx.Slice(qkv, []int32{0, 0, 0}, []int32{B, L, dim})
-	k := mlx.Slice(qkv, []int32{0, 0, dim}, []int32{B, L, 2 * dim})
-	v := mlx.Slice(qkv, []int32{0, 0, 2 * dim}, []int32{B, L, 3 * dim})
-	return q, k, v
-}
-
-// applyQKNorm applies RMSNorm with learned scale (no bias)
-// Uses the optimized mlx_fast_rms_norm
-func applyQKNorm(x *mlx.Array, scale *mlx.Array) *mlx.Array {
-	return mlx.RMSNorm(x, scale, 1e-6)
-}

x/imagegen/models/flux2/vae.go

@@ -1,804 +0,0 @@
-//go:build mlx
-
-package flux2
-
-import (
-	"fmt"
-	"math"
-
-	"github.com/ollama/ollama/x/imagegen"
-	"github.com/ollama/ollama/x/imagegen/mlx"
-	"github.com/ollama/ollama/x/imagegen/nn"
-	"github.com/ollama/ollama/x/imagegen/safetensors"
-	"github.com/ollama/ollama/x/imagegen/vae"
-)
-
-// VAEConfig holds AutoencoderKLFlux2 configuration
-type VAEConfig struct {
-	ActFn             string  `json:"act_fn"`              // "silu"
-	BatchNormEps      float32 `json:"batch_norm_eps"`      // 0.0001
-	BatchNormMomentum float32 `json:"batch_norm_momentum"` // 0.1
-	BlockOutChannels  []int32 `json:"block_out_channels"`  // [128, 256, 512, 512]
-	ForceUpcast       bool    `json:"force_upcast"`        // true
-	InChannels        int32   `json:"in_channels"`         // 3
-	LatentChannels    int32   `json:"latent_channels"`     // 32
-	LayersPerBlock    int32   `json:"layers_per_block"`    // 2
-	MidBlockAddAttn   bool    `json:"mid_block_add_attention"` // true
-	NormNumGroups     int32   `json:"norm_num_groups"`     // 32
-	OutChannels       int32   `json:"out_channels"`        // 3
-	PatchSize         []int32 `json:"patch_size"`          // [2, 2]
-	SampleSize        int32   `json:"sample_size"`         // 1024
-	UsePostQuantConv  bool    `json:"use_post_quant_conv"` // true
-	UseQuantConv      bool    `json:"use_quant_conv"`      // true
-}
-
-// BatchNorm2D implements 2D batch normalization with running statistics
-type BatchNorm2D struct {
-	RunningMean *mlx.Array // [C]
-	RunningVar  *mlx.Array // [C]
-	Weight      *mlx.Array // [C] gamma
-	Bias        *mlx.Array // [C] beta
-	Eps         float32
-	Momentum    float32
-}
-
-// Forward applies batch normalization (inference mode - uses running stats)
-// Input and output are in NHWC format [B, H, W, C]
-func (bn *BatchNorm2D) Forward(x *mlx.Array) *mlx.Array {
-	shape := x.Shape()
-	C := shape[3]
-
-	// Reshape stats for broadcasting [1, 1, 1, C]
-	mean := mlx.Reshape(bn.RunningMean, 1, 1, 1, C)
-	variance := mlx.Reshape(bn.RunningVar, 1, 1, 1, C)
-
-	// Normalize: (x - mean) / sqrt(var + eps)
-	xNorm := mlx.Sub(x, mean)
-	xNorm = mlx.Div(xNorm, mlx.Sqrt(mlx.AddScalar(variance, bn.Eps)))
-
-	// Scale and shift (only if affine=True)
-	if bn.Weight != nil {
-		weight := mlx.Reshape(bn.Weight, 1, 1, 1, C)
-		xNorm = mlx.Mul(xNorm, weight)
-	}
-	if bn.Bias != nil {
-		bias := mlx.Reshape(bn.Bias, 1, 1, 1, C)
-		xNorm = mlx.Add(xNorm, bias)
-	}
-
-	return xNorm
-}
-
-// Denormalize inverts the batch normalization
-// Used when decoding latents
-func (bn *BatchNorm2D) Denormalize(x *mlx.Array) *mlx.Array {
-	shape := x.Shape()
-	C := shape[3]
-
-	// Reshape stats for broadcasting [1, 1, 1, C]
-	mean := mlx.Reshape(bn.RunningMean, 1, 1, 1, C)
-	variance := mlx.Reshape(bn.RunningVar, 1, 1, 1, C)
-
-	// Inverse: first undo affine, then undo normalization
-	// For affine=False: x_denorm = x * sqrt(var + eps) + mean
-	if bn.Bias != nil {
-		bias := mlx.Reshape(bn.Bias, 1, 1, 1, C)
-		x = mlx.Sub(x, bias)
-	}
-	if bn.Weight != nil {
-		weight := mlx.Reshape(bn.Weight, 1, 1, 1, C)
-		x = mlx.Div(x, weight)
-	}
-	x = mlx.Mul(x, mlx.Sqrt(mlx.AddScalar(variance, bn.Eps)))
-	x = mlx.Add(x, mean)
-
-	return x
-}
-
-// GroupNormLayer implements group normalization
-// Reused from zimage package pattern
-type GroupNormLayer struct {
-	Weight    *mlx.Array `weight:"weight"`
-	Bias      *mlx.Array `weight:"bias"`
-	NumGroups int32
-	Eps       float32
-}
-
-// Forward applies group normalization
-// Input and output are in NHWC format [B, H, W, C]
-func (gn *GroupNormLayer) Forward(x *mlx.Array) *mlx.Array {
-	shape := x.Shape()
-	B := shape[0]
-	H := shape[1]
-	W := shape[2]
-	C := shape[3]
-
-	// Reshape to [B, H, W, groups, C/groups]
-	groupSize := C / gn.NumGroups
-	x = mlx.Reshape(x, B, H, W, gn.NumGroups, groupSize)
-
-	// Compute mean and variance per group
-	mean := mlx.Mean(x, 1, true)
-	mean = mlx.Mean(mean, 2, true)
-	mean = mlx.Mean(mean, 4, true)
-
-	xCentered := mlx.Sub(x, mean)
-
-	sq := mlx.Square(xCentered)
-	variance := mlx.Mean(sq, 1, true)
-	variance = mlx.Mean(variance, 2, true)
-	variance = mlx.Mean(variance, 4, true)
-
-	// Normalize
-	xNorm := mlx.Div(xCentered, mlx.Sqrt(mlx.AddScalar(variance, gn.Eps)))
-
-	// Reshape back to [B, H, W, C]
-	xNorm = mlx.Reshape(xNorm, B, H, W, C)
-
-	// Scale and shift
-	if gn.Weight != nil {
-		weight := mlx.Reshape(gn.Weight, 1, 1, 1, C)
-		xNorm = mlx.Mul(xNorm, weight)
-	}
-	if gn.Bias != nil {
-		bias := mlx.Reshape(gn.Bias, 1, 1, 1, C)
-		xNorm = mlx.Add(xNorm, bias)
-	}
-
-	return xNorm
-}
-
-// Conv2D represents a 2D convolution layer (reused pattern)
-type Conv2D struct {
-	Weight  *mlx.Array `weight:"weight"`
-	Bias    *mlx.Array `weight:"bias,optional"`
-	Stride  int32
-	Padding int32
-}
-
-// Transform implements safetensors.Transformer to transpose weights from PyTorch's OIHW to MLX's OHWI.
-func (conv *Conv2D) Transform(field string, arr *mlx.Array) *mlx.Array {
-	if field == "Weight" {
-		return mlx.Transpose(arr, 0, 2, 3, 1)
-	}
-	return arr
-}
-
-// Forward applies convolution (NHWC format)
-func (conv *Conv2D) Forward(x *mlx.Array) *mlx.Array {
-	out := mlx.Conv2d(x, conv.Weight, conv.Stride, conv.Padding)
-
-	if conv.Bias != nil {
-		bias := mlx.Reshape(conv.Bias, 1, 1, 1, conv.Bias.Dim(0))
-		out = mlx.Add(out, bias)
-	}
-
-	return out
-}
-
-// ResnetBlock2D implements a ResNet block for VAE
-type ResnetBlock2D struct {
-	Norm1        *GroupNormLayer `weight:"norm1"`
-	Conv1        *Conv2D         `weight:"conv1"`
-	Norm2        *GroupNormLayer `weight:"norm2"`
-	Conv2        *Conv2D         `weight:"conv2"`
-	ConvShortcut *Conv2D         `weight:"conv_shortcut,optional"`
-}
-
-// Forward applies the ResNet block
-func (rb *ResnetBlock2D) Forward(x *mlx.Array) *mlx.Array {
-	h := rb.Norm1.Forward(x)
-	h = mlx.SiLU(h)
-	h = rb.Conv1.Forward(h)
-
-	h = rb.Norm2.Forward(h)
-	h = mlx.SiLU(h)
-	h = rb.Conv2.Forward(h)
-
-	if rb.ConvShortcut != nil {
-		x = rb.ConvShortcut.Forward(x)
-	}
-
-	return mlx.Add(h, x)
-}
-
-// VAEAttentionBlock implements self-attention for VAE
-type VAEAttentionBlock struct {
-	GroupNorm *GroupNormLayer `weight:"group_norm"`
-	ToQ       nn.LinearLayer  `weight:"to_q"`
-	ToK       nn.LinearLayer  `weight:"to_k"`
-	ToV       nn.LinearLayer  `weight:"to_v"`
-	ToOut     nn.LinearLayer  `weight:"to_out.0"`
-}
-
-// Forward applies attention (NHWC format)
-func (ab *VAEAttentionBlock) Forward(x *mlx.Array) *mlx.Array {
-	residual := x
-	shape := x.Shape()
-	B := shape[0]
-	H := shape[1]
-	W := shape[2]
-	C := shape[3]
-
-	h := ab.GroupNorm.Forward(x)
-	h = mlx.Reshape(h, B, H*W, C)
-
-	q := ab.ToQ.Forward(h)
-	k := ab.ToK.Forward(h)
-	v := ab.ToV.Forward(h)
-
-	q = mlx.ExpandDims(q, 1)
-	k = mlx.ExpandDims(k, 1)
-	v = mlx.ExpandDims(v, 1)
-
-	scale := float32(1.0 / math.Sqrt(float64(C)))
-	out := mlx.ScaledDotProductAttention(q, k, v, scale, false)
-	out = mlx.Squeeze(out, 1)
-
-	out = ab.ToOut.Forward(out)
-	out = mlx.Reshape(out, B, H, W, C)
-	out = mlx.Add(out, residual)
-
-	return out
-}
-
-// UpDecoderBlock2D implements an upsampling decoder block
-type UpDecoderBlock2D struct {
-	ResnetBlocks []*ResnetBlock2D
-	Upsample     *Conv2D
-}
-
-// Forward applies the up decoder block
-func (ub *UpDecoderBlock2D) Forward(x *mlx.Array) *mlx.Array {
-	for _, resnet := range ub.ResnetBlocks {
-		x = resnet.Forward(x)
-	}
-
-	if ub.Upsample != nil {
-		x = upsample2x(x)
-		x = ub.Upsample.Forward(x)
-	}
-
-	return x
-}
-
-// upsample2x performs 2x nearest neighbor upsampling
-func upsample2x(x *mlx.Array) *mlx.Array {
-	shape := x.Shape()
-	H := shape[1]
-	W := shape[2]
-
-	hIdx := mlx.ArangeInt(0, H, 1, mlx.DtypeInt32)
-	hIdx = mlx.Reshape(hIdx, H, 1)
-	hIdx = mlx.BroadcastTo(hIdx, []int32{H, 2})
-	hIdx = mlx.Reshape(hIdx, H*2)
-
-	wIdx := mlx.ArangeInt(0, W, 1, mlx.DtypeInt32)
-	wIdx = mlx.Reshape(wIdx, W, 1)
-	wIdx = mlx.BroadcastTo(wIdx, []int32{W, 2})
-	wIdx = mlx.Reshape(wIdx, W*2)
-
-	x = mlx.Take(x, hIdx, 1)
-	x = mlx.Take(x, wIdx, 2)
-
-	return x
-}
-
-// VAEMidBlock is the middle block with attention
-type VAEMidBlock struct {
-	Resnet1   *ResnetBlock2D
-	Attention *VAEAttentionBlock
-	Resnet2   *ResnetBlock2D
-}
-
-// Forward applies the mid block
-func (mb *VAEMidBlock) Forward(x *mlx.Array) *mlx.Array {
-	x = mb.Resnet1.Forward(x)
-	x = mb.Attention.Forward(x)
-	x = mb.Resnet2.Forward(x)
-	return x
-}
-
-// DefaultTilingConfig returns reasonable defaults for tiled decoding
-// Matches diffusers: tile_latent_min_size=64, tile_overlap_factor=0.25
-func DefaultTilingConfig() *vae.TilingConfig {
-	return vae.DefaultTilingConfig()
-}
-
-// AutoencoderKLFlux2 is the Flux2 VAE with BatchNorm
-type AutoencoderKLFlux2 struct {
-	Config *VAEConfig
-
-	// Encoder components (for image editing)
-	EncoderConvIn  *Conv2D
-	EncoderMid     *VAEMidBlock
-	EncoderDown    []*DownEncoderBlock2D
-	EncoderNormOut *GroupNormLayer
-	EncoderConvOut *Conv2D
-
-	// Decoder components
-	DecoderConvIn  *Conv2D
-	DecoderMid     *VAEMidBlock
-	DecoderUp      []*UpDecoderBlock2D
-	DecoderNormOut *GroupNormLayer
-	DecoderConvOut *Conv2D
-
-	// Quant conv layers
-	QuantConv     *Conv2D
-	PostQuantConv *Conv2D
-
-	// BatchNorm for latent normalization
-	LatentBN *BatchNorm2D
-
-	// Tiling configuration (nil = no tiling)
-	Tiling *vae.TilingConfig
-}
-
-// DownEncoderBlock2D implements a downsampling encoder block
-type DownEncoderBlock2D struct {
-	ResnetBlocks []*ResnetBlock2D
-	Downsample   *Conv2D
-}
-
-// Forward applies the down encoder block
-func (db *DownEncoderBlock2D) Forward(x *mlx.Array) *mlx.Array {
-	for _, resnet := range db.ResnetBlocks {
-		x = resnet.Forward(x)
-	}
-
-	if db.Downsample != nil {
-		// Pad then conv with stride 2
-		x = mlx.Pad(x, []int32{0, 0, 0, 1, 0, 1, 0, 0})
-		x = db.Downsample.Forward(x)
-	}
-
-	return x
-}
-
-// Load loads the Flux2 VAE from ollama blob storage.
-func (m *AutoencoderKLFlux2) Load(manifest *imagegen.ModelManifest) error {
-	fmt.Print("  Loading VAE... ")
-
-	// Load config from blob
-	var cfg VAEConfig
-	if err := manifest.ReadConfigJSON("vae/config.json", &cfg); err != nil {
-		return fmt.Errorf("config: %w", err)
-	}
-	m.Config = &cfg
-
-	// Load weights from tensor blobs
-	weights, err := imagegen.LoadWeightsFromManifest(manifest, "vae")
-	if err != nil {
-		return fmt.Errorf("weights: %w", err)
-	}
-	if err := weights.Load(0); err != nil {
-		return fmt.Errorf("load weights: %w", err)
-	}
-	defer weights.ReleaseAll()
-
-	return m.loadWeights(weights, &cfg)
-}
-
-// loadWeights loads VAE weights from any WeightSource
-func (m *AutoencoderKLFlux2) loadWeights(weights safetensors.WeightSource, cfg *VAEConfig) error {
-	var err error
-
-	// Load encoder components (for image conditioning)
-	if err := m.loadEncoderWeights(weights, cfg); err != nil {
-		return fmt.Errorf("encoder: %w", err)
-	}
-
-	// Load decoder conv_in
-	m.DecoderConvIn = &Conv2D{Stride: 1, Padding: 1}
-	if err := safetensors.LoadModule(m.DecoderConvIn, weights, "decoder.conv_in"); err != nil {
-		return fmt.Errorf("decoder.conv_in: %w", err)
-	}
-
-	// Load mid block
-	m.DecoderMid, err = loadVAEMidBlock(weights, "decoder.mid_block", cfg.NormNumGroups)
-	if err != nil {
-		return fmt.Errorf("decoder.mid_block: %w", err)
-	}
-
-	// Load up blocks
-	numBlocks := len(cfg.BlockOutChannels)
-	m.DecoderUp = make([]*UpDecoderBlock2D, numBlocks)
-	for i := 0; i < numBlocks; i++ {
-		prefix := fmt.Sprintf("decoder.up_blocks.%d", i)
-		hasUpsample := i < numBlocks-1
-		m.DecoderUp[i], err = loadUpDecoderBlock2D(weights, prefix, cfg.LayersPerBlock+1, cfg.NormNumGroups, hasUpsample)
-		if err != nil {
-			return fmt.Errorf("%s: %w", prefix, err)
-		}
-	}
-
-	// Load decoder conv_norm_out and conv_out
-	m.DecoderNormOut = &GroupNormLayer{NumGroups: cfg.NormNumGroups, Eps: 1e-5}
-	if err := safetensors.LoadModule(m.DecoderNormOut, weights, "decoder.conv_norm_out"); err != nil {
-		return fmt.Errorf("decoder.conv_norm_out: %w", err)
-	}
-
-	m.DecoderConvOut = &Conv2D{Stride: 1, Padding: 1}
-	if err := safetensors.LoadModule(m.DecoderConvOut, weights, "decoder.conv_out"); err != nil {
-		return fmt.Errorf("decoder.conv_out: %w", err)
-	}
-
-	// Load post_quant_conv
-	if cfg.UsePostQuantConv {
-		m.PostQuantConv = &Conv2D{Stride: 1, Padding: 0}
-		if err := safetensors.LoadModule(m.PostQuantConv, weights, "post_quant_conv"); err != nil {
-			return fmt.Errorf("post_quant_conv: %w", err)
-		}
-	}
-
-	// Load latent BatchNorm (affine=False, so no weight/bias)
-	bnMean, err := weights.GetTensor("bn.running_mean")
-	if err != nil {
-		return fmt.Errorf("bn.running_mean: %w", err)
-	}
-	bnVar, err := weights.GetTensor("bn.running_var")
-	if err != nil {
-		return fmt.Errorf("bn.running_var: %w", err)
-	}
-	m.LatentBN = &BatchNorm2D{
-		RunningMean: bnMean,
-		RunningVar:  bnVar,
-		Weight:      nil, // affine=False
-		Bias:        nil, // affine=False
-		Eps:         cfg.BatchNormEps,
-		Momentum:    cfg.BatchNormMomentum,
-	}
-
-	fmt.Println("✓")
-	return nil
-}
-
-// loadVAEMidBlock loads the mid block.
-func loadVAEMidBlock(weights safetensors.WeightSource, prefix string, numGroups int32) (*VAEMidBlock, error) {
-	resnet1, err := loadResnetBlock2D(weights, prefix+".resnets.0", numGroups)
-	if err != nil {
-		return nil, err
-	}
-
-	attention, err := loadVAEAttentionBlock(weights, prefix+".attentions.0", numGroups)
-	if err != nil {
-		return nil, err
-	}
-
-	resnet2, err := loadResnetBlock2D(weights, prefix+".resnets.1", numGroups)
-	if err != nil {
-		return nil, err
-	}
-
-	return &VAEMidBlock{
-		Resnet1:   resnet1,
-		Attention: attention,
-		Resnet2:   resnet2,
-	}, nil
-}
-
-// loadResnetBlock2D loads a ResNet block.
-func loadResnetBlock2D(weights safetensors.WeightSource, prefix string, numGroups int32) (*ResnetBlock2D, error) {
-	block := &ResnetBlock2D{
-		Norm1:        &GroupNormLayer{NumGroups: numGroups, Eps: 1e-5},
-		Conv1:        &Conv2D{Stride: 1, Padding: 1},
-		Norm2:        &GroupNormLayer{NumGroups: numGroups, Eps: 1e-5},
-		Conv2:        &Conv2D{Stride: 1, Padding: 1},
-		ConvShortcut: &Conv2D{Stride: 1, Padding: 0}, // Pre-allocate for optional loading
-	}
-	if err := safetensors.LoadModule(block, weights, prefix); err != nil {
-		return nil, err
-	}
-	// If ConvShortcut wasn't loaded (no weights found), nil it out
-	if block.ConvShortcut.Weight == nil {
-		block.ConvShortcut = nil
-	}
-	return block, nil
-}
-
-// loadVAEAttentionBlock loads an attention block using LoadModule.
-func loadVAEAttentionBlock(weights safetensors.WeightSource, prefix string, numGroups int32) (*VAEAttentionBlock, error) {
-	ab := &VAEAttentionBlock{
-		GroupNorm: &GroupNormLayer{NumGroups: numGroups, Eps: 1e-5},
-	}
-	if err := safetensors.LoadModule(ab, weights, prefix); err != nil {
-		return nil, err
-	}
-	return ab, nil
-}
-
-// loadUpDecoderBlock2D loads an up decoder block.
-func loadUpDecoderBlock2D(weights safetensors.WeightSource, prefix string, numLayers, numGroups int32, hasUpsample bool) (*UpDecoderBlock2D, error) {
-	resnets := make([]*ResnetBlock2D, numLayers)
-	for i := int32(0); i < numLayers; i++ {
-		resPrefix := fmt.Sprintf("%s.resnets.%d", prefix, i)
-		resnet, err := loadResnetBlock2D(weights, resPrefix, numGroups)
-		if err != nil {
-			return nil, err
-		}
-		resnets[i] = resnet
-	}
-
-	var upsample *Conv2D
-	if hasUpsample {
-		upsample = &Conv2D{Stride: 1, Padding: 1}
-		if err := safetensors.LoadModule(upsample, weights, prefix+".upsamplers.0.conv"); err != nil {
-			return nil, err
-		}
-	}
-
-	return &UpDecoderBlock2D{
-		ResnetBlocks: resnets,
-		Upsample:     upsample,
-	}, nil
-}
-
-// Patchify converts latents [B, C, H, W] to patches [B, H*W/4, C*4] using 2x2 patches
-// This is the inverse of the VAE's patchify for feeding to transformer
-func (vae *AutoencoderKLFlux2) Patchify(latents *mlx.Array) *mlx.Array {
-	shape := latents.Shape()
-	B := shape[0]
-	C := shape[1]
-	H := shape[2]
-	W := shape[3]
-
-	patchH := vae.Config.PatchSize[0]
-	patchW := vae.Config.PatchSize[1]
-
-	pH := H / patchH
-	pW := W / patchW
-
-	// [B, C, H, W] -> [B, C, pH, patchH, pW, patchW]
-	x := mlx.Reshape(latents, B, C, pH, patchH, pW, patchW)
-	// [B, C, pH, patchH, pW, patchW] -> [B, pH, pW, C, patchH, patchW]
-	x = mlx.Transpose(x, 0, 2, 4, 1, 3, 5)
-	// [B, pH, pW, C, patchH, patchW] -> [B, pH*pW, C*patchH*patchW]
-	return mlx.Reshape(x, B, pH*pW, C*patchH*patchW)
-}
-
-// Unpatchify converts patches [B, L, C*4] back to [B, C, H, W]
-func (vae *AutoencoderKLFlux2) Unpatchify(patches *mlx.Array, pH, pW, C int32) *mlx.Array {
-	shape := patches.Shape()
-	B := shape[0]
-
-	patchH := vae.Config.PatchSize[0]
-	patchW := vae.Config.PatchSize[1]
-
-	// [B, pH*pW, C*patchH*patchW] -> [B, pH, pW, C, patchH, patchW]
-	x := mlx.Reshape(patches, B, pH, pW, C, patchH, patchW)
-	// [B, pH, pW, C, patchH, patchW] -> [B, C, pH, patchH, pW, patchW]
-	x = mlx.Transpose(x, 0, 3, 1, 4, 2, 5)
-	// [B, C, pH, patchH, pW, patchW] -> [B, C, H, W]
-	H := pH * patchH
-	W := pW * patchW
-	return mlx.Reshape(x, B, C, H, W)
-}
-
-// denormalizePatchified applies inverse batch normalization to patchified latents.
-// Input: [B, L, 128] where 128 = 32 latent channels * 4 (2x2 patch)
-// Output: [B, L, 128] denormalized
-func (vae *AutoencoderKLFlux2) denormalizePatchified(x *mlx.Array) *mlx.Array {
-	shape := x.Shape()
-	C := shape[2] // 128
-
-	// Reshape stats for broadcasting [1, 1, C]
-	mean := mlx.Reshape(vae.LatentBN.RunningMean, 1, 1, C)
-	variance := mlx.Reshape(vae.LatentBN.RunningVar, 1, 1, C)
-
-	// Inverse BN (affine=False): x_denorm = x * sqrt(var + eps) + mean
-	if vae.LatentBN.Bias != nil {
-		bias := mlx.Reshape(vae.LatentBN.Bias, 1, 1, C)
-		x = mlx.Sub(x, bias)
-	}
-	if vae.LatentBN.Weight != nil {
-		weight := mlx.Reshape(vae.LatentBN.Weight, 1, 1, C)
-		x = mlx.Div(x, weight)
-	}
-	x = mlx.Mul(x, mlx.Sqrt(mlx.AddScalar(variance, vae.LatentBN.Eps)))
-	x = mlx.Add(x, mean)
-
-	return x
-}
-
-// Decode decodes latent patches to images.
-// If Tiling is set, uses tiled decoding to reduce memory for large images.
-// latents: [B, L, C*4] patchified latents from transformer
-// pH, pW: patch grid dimensions
-// Returns: [B, 3, H, W] image tensor
-func (v *AutoencoderKLFlux2) Decode(latents *mlx.Array, pH, pW int32) *mlx.Array {
-	// Denormalize patchified latents
-	z := v.denormalizePatchified(latents)
-
-	// Unpatchify: [B, L, C*4] -> [B, C, H, W]
-	z = v.Unpatchify(z, pH, pW, v.Config.LatentChannels)
-
-	// Convert NCHW -> NHWC for processing
-	z = mlx.Transpose(z, 0, 2, 3, 1)
-
-	// Use tiled decoding if enabled
-	if v.Tiling != nil {
-		mlx.Eval(z)
-		return vae.DecodeTiled(z, v.Tiling, v.decodeTile)
-	}
-
-	// Direct decode (no tiling)
-	h := v.decodeTile(z)
-	h = mlx.ClipScalar(h, 0.0, 1.0, true, true)
-	h = mlx.Transpose(h, 0, 3, 1, 2)
-	return h
-}
-
-// decodeTile decodes a single latent tile to pixels (internal helper)
-// z: [B, H, W, C] latent tile in NHWC format
-// Returns: [B, H*8, W*8, 3] pixel tile in NHWC format (before clipping)
-func (vae *AutoencoderKLFlux2) decodeTile(z *mlx.Array) *mlx.Array {
-	// Post-quant conv
-	if vae.PostQuantConv != nil {
-		z = vae.PostQuantConv.Forward(z)
-	}
-
-	// Decoder
-	h := vae.DecoderConvIn.Forward(z)
-	h = vae.DecoderMid.Forward(h)
-
-	for _, upBlock := range vae.DecoderUp {
-		h = upBlock.Forward(h)
-	}
-
-	h = vae.DecoderNormOut.Forward(h)
-	h = mlx.SiLU(h)
-	h = vae.DecoderConvOut.Forward(h)
-
-	// VAE outputs [-1, 1], convert to [0, 1]
-	h = mlx.MulScalar(h, 0.5)
-	h = mlx.AddScalar(h, 0.5)
-
-	return h
-}
-
-// loadEncoderWeights loads the encoder components for image conditioning
-func (m *AutoencoderKLFlux2) loadEncoderWeights(weights safetensors.WeightSource, cfg *VAEConfig) error {
-	var err error
-
-	// Load encoder conv_in
-	m.EncoderConvIn = &Conv2D{Stride: 1, Padding: 1}
-	if err := safetensors.LoadModule(m.EncoderConvIn, weights, "encoder.conv_in"); err != nil {
-		return fmt.Errorf("encoder.conv_in: %w", err)
-	}
-
-	// Load encoder down blocks
-	numBlocks := len(cfg.BlockOutChannels)
-	m.EncoderDown = make([]*DownEncoderBlock2D, numBlocks)
-	for i := 0; i < numBlocks; i++ {
-		prefix := fmt.Sprintf("encoder.down_blocks.%d", i)
-		hasDownsample := i < numBlocks-1
-		m.EncoderDown[i], err = loadDownEncoderBlock2D(weights, prefix, cfg.LayersPerBlock, cfg.NormNumGroups, hasDownsample)
-		if err != nil {
-			return fmt.Errorf("%s: %w", prefix, err)
-		}
-	}
-
-	// Load encoder mid block
-	m.EncoderMid, err = loadVAEMidBlock(weights, "encoder.mid_block", cfg.NormNumGroups)
-	if err != nil {
-		return fmt.Errorf("encoder.mid_block: %w", err)
-	}
-
-	// Load encoder conv_norm_out and conv_out
-	m.EncoderNormOut = &GroupNormLayer{NumGroups: cfg.NormNumGroups, Eps: 1e-5}
-	if err := safetensors.LoadModule(m.EncoderNormOut, weights, "encoder.conv_norm_out"); err != nil {
-		return fmt.Errorf("encoder.conv_norm_out: %w", err)
-	}
-
-	m.EncoderConvOut = &Conv2D{Stride: 1, Padding: 1}
-	if err := safetensors.LoadModule(m.EncoderConvOut, weights, "encoder.conv_out"); err != nil {
-		return fmt.Errorf("encoder.conv_out: %w", err)
-	}
-
-	// Load quant_conv (for encoding)
-	if cfg.UseQuantConv {
-		m.QuantConv = &Conv2D{Stride: 1, Padding: 0}
-		if err := safetensors.LoadModule(m.QuantConv, weights, "quant_conv"); err != nil {
-			return fmt.Errorf("quant_conv: %w", err)
-		}
-	}
-
-	return nil
-}
-
-// loadDownEncoderBlock2D loads a down encoder block.
-func loadDownEncoderBlock2D(weights safetensors.WeightSource, prefix string, numLayers, numGroups int32, hasDownsample bool) (*DownEncoderBlock2D, error) {
-	resnets := make([]*ResnetBlock2D, numLayers)
-	for i := int32(0); i < numLayers; i++ {
-		resPrefix := fmt.Sprintf("%s.resnets.%d", prefix, i)
-		resnet, err := loadResnetBlock2D(weights, resPrefix, numGroups)
-		if err != nil {
-			return nil, err
-		}
-		resnets[i] = resnet
-	}
-
-	var downsample *Conv2D
-	if hasDownsample {
-		downsample = &Conv2D{Stride: 2, Padding: 0}
-		if err := safetensors.LoadModule(downsample, weights, prefix+".downsamplers.0.conv"); err != nil {
-			return nil, err
-		}
-	}
-
-	return &DownEncoderBlock2D{
-		ResnetBlocks: resnets,
-		Downsample:   downsample,
-	}, nil
-}
-
-// EncodeImage encodes an image to normalized latents.
-// image: [B, 3, H, W] image tensor in [-1, 1]
-// Returns: [B, L, C*4] patchified normalized latents
-func (vae *AutoencoderKLFlux2) EncodeImage(image *mlx.Array) *mlx.Array {
-	// Convert NCHW -> NHWC
-	x := mlx.Transpose(image, 0, 2, 3, 1)
-
-	// Encoder
-	h := vae.EncoderConvIn.Forward(x)
-
-	for _, downBlock := range vae.EncoderDown {
-		h = downBlock.Forward(h)
-	}
-
-	h = vae.EncoderMid.Forward(h)
-	h = vae.EncoderNormOut.Forward(h)
-	h = mlx.SiLU(h)
-	h = vae.EncoderConvOut.Forward(h)
-
-	// Quant conv outputs [B, H, W, 2*latent_channels] (mean + logvar)
-	if vae.QuantConv != nil {
-		h = vae.QuantConv.Forward(h)
-	}
-
-	// Take only the mean (first latent_channels) - deterministic encoding
-	// h is [B, H, W, 64] -> take first 32 channels for mean
-	shape := h.Shape()
-	latentChannels := vae.Config.LatentChannels // 32
-	h = mlx.Slice(h, []int32{0, 0, 0, 0}, []int32{shape[0], shape[1], shape[2], latentChannels})
-
-	// Convert NHWC -> NCHW for patchifying
-	h = mlx.Transpose(h, 0, 3, 1, 2)
-
-	// Patchify: [B, C, H, W] -> [B, L, C*4]
-	h = vae.Patchify(h)
-
-	// Apply BatchNorm on patchified latents [B, L, 128]
-	// The BatchNorm has 128 channels matching the patchified dimension
-	h = vae.normalizePatchified(h)
-
-	return h
-}
-
-// normalizePatchified applies batch normalization to patchified latents.
-// Input: [B, L, 128] where 128 = 32 latent channels * 4 (2x2 patch)
-// Output: [B, L, 128] normalized
-func (vae *AutoencoderKLFlux2) normalizePatchified(x *mlx.Array) *mlx.Array {
-	shape := x.Shape()
-	C := shape[2] // 128
-
-	// Reshape stats for broadcasting [1, 1, C]
-	mean := mlx.Reshape(vae.LatentBN.RunningMean, 1, 1, C)
-	variance := mlx.Reshape(vae.LatentBN.RunningVar, 1, 1, C)
-
-	// Normalize: (x - mean) / sqrt(var + eps)
-	xNorm := mlx.Sub(x, mean)
-	xNorm = mlx.Div(xNorm, mlx.Sqrt(mlx.AddScalar(variance, vae.LatentBN.Eps)))
-
-	// Scale and shift (only if affine=True)
-	if vae.LatentBN.Weight != nil {
-		weight := mlx.Reshape(vae.LatentBN.Weight, 1, 1, C)
-		xNorm = mlx.Mul(xNorm, weight)
-	}
-	if vae.LatentBN.Bias != nil {
-		bias := mlx.Reshape(vae.LatentBN.Bias, 1, 1, C)
-		xNorm = mlx.Add(xNorm, bias)
-	}
-
-	return xNorm
-}

x/imagegen/models/qwen3/text_encoder.go

@@ -1,390 +0,0 @@
-//go:build mlx
-
-// Package qwen3 provides a shared Qwen3 text encoder used by multiple image generation models.
-package qwen3
-
-import (
-	"fmt"
-	"math"
-
-	"github.com/ollama/ollama/x/imagegen"
-	"github.com/ollama/ollama/x/imagegen/mlx"
-	"github.com/ollama/ollama/x/imagegen/nn"
-	"github.com/ollama/ollama/x/imagegen/safetensors"
-	"github.com/ollama/ollama/x/imagegen/tokenizer"
-)
-
-// Config holds Qwen3 text encoder configuration
-type Config struct {
-	HiddenSize        int32   `json:"hidden_size"`
-	NumHiddenLayers   int32   `json:"num_hidden_layers"`
-	IntermediateSize  int32   `json:"intermediate_size"`
-	NumAttentionHeads int32   `json:"num_attention_heads"`
-	NumKeyValueHeads  int32   `json:"num_key_value_heads"`
-	VocabSize         int32   `json:"vocab_size"`
-	RMSNormEps        float32 `json:"rms_norm_eps"`
-	RopeTheta         float32 `json:"rope_theta"`
-	HeadDim           int32   `json:"head_dim"`
-}
-
-// Attention implements Qwen3 attention with QK norms
-type Attention struct {
-	QProj nn.LinearLayer `weight:"q_proj"`
-	KProj nn.LinearLayer `weight:"k_proj"`
-	VProj nn.LinearLayer `weight:"v_proj"`
-	OProj nn.LinearLayer `weight:"o_proj"`
-	QNorm *nn.RMSNorm    `weight:"q_norm"`
-	KNorm *nn.RMSNorm    `weight:"k_norm"`
-	// Computed fields
-	NHeads    int32
-	NKVHeads  int32
-	HeadDim   int32
-	Scale     float32
-	RopeTheta float32
-}
-
-// applyRoPEQwen3 applies the custom RoPE for Qwen3 text encoder
-func applyRoPEQwen3(x *mlx.Array, seqLen int32, theta float32) *mlx.Array {
-	shape := x.Shape()
-	B := shape[0]
-	L := shape[1]
-	H := shape[2]
-	D := shape[3]
-	half := D / 2
-
-	freqsArr := make([]float32, half)
-	logTheta := float32(math.Log(float64(theta)))
-	for i := int32(0); i < half; i++ {
-		freqsArr[i] = float32(math.Exp(float64(-logTheta * float32(i) / float32(half))))
-	}
-	freqs := mlx.NewArray(freqsArr, []int32{half})
-
-	posArr := make([]float32, seqLen)
-	for i := int32(0); i < seqLen; i++ {
-		posArr[i] = float32(i)
-	}
-	pos := mlx.NewArray(posArr, []int32{seqLen})
-
-	posExpanded := mlx.Reshape(pos, seqLen, 1)
-	freqsExpanded := mlx.Reshape(freqs, 1, half)
-	args := mlx.Mul(posExpanded, freqsExpanded)
-
-	cosVals := mlx.Cos(args)
-	sinVals := mlx.Sin(args)
-	cosVals = mlx.Reshape(cosVals, seqLen, 1, half)
-	sinVals = mlx.Reshape(sinVals, seqLen, 1, half)
-
-	x1 := mlx.Slice(x, []int32{0, 0, 0, 0}, []int32{B, L, H, half})
-	x2 := mlx.Slice(x, []int32{0, 0, 0, half}, []int32{B, L, H, D})
-
-	part1 := mlx.Sub(mlx.Mul(x1, cosVals), mlx.Mul(x2, sinVals))
-	part2 := mlx.Add(mlx.Mul(x1, sinVals), mlx.Mul(x2, cosVals))
-
-	return mlx.Concatenate([]*mlx.Array{part1, part2}, 3)
-}
-
-// Forward computes attention with causal masking and optional padding mask
-func (attn *Attention) Forward(x *mlx.Array, mask *mlx.Array, maskMode string) *mlx.Array {
-	shape := x.Shape()
-	B := shape[0]
-	L := shape[1]
-
-	q := attn.QProj.Forward(x)
-	k := attn.KProj.Forward(x)
-	v := attn.VProj.Forward(x)
-
-	q = mlx.Reshape(q, B, L, attn.NHeads, attn.HeadDim)
-	k = mlx.Reshape(k, B, L, attn.NKVHeads, attn.HeadDim)
-	v = mlx.Reshape(v, B, L, attn.NKVHeads, attn.HeadDim)
-
-	// QK norm uses 1e-6 hardcoded (Qwen3 specific)
-	q = attn.QNorm.Forward(q, 1e-6)
-	k = attn.KNorm.Forward(k, 1e-6)
-
-	q = applyRoPEQwen3(q, L, attn.RopeTheta)
-	k = applyRoPEQwen3(k, L, attn.RopeTheta)
-
-	q = mlx.Transpose(q, 0, 2, 1, 3)
-	k = mlx.Transpose(k, 0, 2, 1, 3)
-	v = mlx.Transpose(v, 0, 2, 1, 3)
-
-	if attn.NKVHeads < attn.NHeads {
-		repeats := attn.NHeads / attn.NKVHeads
-		k = repeatKV(k, repeats)
-		v = repeatKV(v, repeats)
-	}
-
-	out := mlx.ScaledDotProductAttentionWithSinks(q, k, v, attn.Scale, maskMode, mask, nil)
-
-	out = mlx.Transpose(out, 0, 2, 1, 3)
-	out = mlx.Reshape(out, B, L, attn.NHeads*attn.HeadDim)
-
-	out = attn.OProj.Forward(out)
-
-	return out
-}
-
-// repeatKV repeats key/value heads for GQA
-func repeatKV(x *mlx.Array, repeats int32) *mlx.Array {
-	if repeats == 1 {
-		return x
-	}
-	shape := x.Shape()
-	x = mlx.ExpandDims(x, 2)
-	x = mlx.Tile(x, []int32{1, 1, repeats, 1, 1})
-	return mlx.Reshape(x, shape[0], shape[1]*repeats, shape[2], shape[3])
-}
-
-// MLP implements Qwen3 SwiGLU MLP
-type MLP struct {
-	GateProj nn.LinearLayer `weight:"gate_proj"`
-	UpProj   nn.LinearLayer `weight:"up_proj"`
-	DownProj nn.LinearLayer `weight:"down_proj"`
-}
-
-// Forward applies the MLP
-func (m *MLP) Forward(x *mlx.Array) *mlx.Array {
-	gate := m.GateProj.Forward(x)
-	gate = mlx.SiLU(gate)
-	up := m.UpProj.Forward(x)
-	h := mlx.Mul(gate, up)
-	return m.DownProj.Forward(h)
-}
-
-// Block represents a single Qwen3 transformer block
-type Block struct {
-	Attention         *Attention  `weight:"self_attn"`
-	MLP               *MLP        `weight:"mlp"`
-	InputLayerNorm    *nn.RMSNorm `weight:"input_layernorm"`
-	PostAttnLayerNorm *nn.RMSNorm `weight:"post_attention_layernorm"`
-}
-
-// Forward applies the Qwen3 block
-func (qb *Block) Forward(x *mlx.Array, eps float32, mask *mlx.Array, maskMode string) *mlx.Array {
-	h := qb.InputLayerNorm.Forward(x, eps)
-	attnOut := qb.Attention.Forward(h, mask, maskMode)
-	x = mlx.Add(x, attnOut)
-
-	h = qb.PostAttnLayerNorm.Forward(x, eps)
-	mlpOut := qb.MLP.Forward(h)
-	x = mlx.Add(x, mlpOut)
-
-	return x
-}
-
-// TextEncoder is the full Qwen3 encoder
-type TextEncoder struct {
-	EmbedTokens *nn.Embedding `weight:"model.embed_tokens"`
-	Layers      []*Block      `weight:"model.layers"`
-	FinalNorm   *nn.RMSNorm   `weight:"model.norm"`
-	*Config
-}
-
-// Load loads the Qwen3 text encoder from ollama blob storage.
-func (m *TextEncoder) Load(manifest *imagegen.ModelManifest, configPath string) error {
-	fmt.Print("  Loading text encoder... ")
-
-	// Load config from blob
-	var cfg Config
-	if err := manifest.ReadConfigJSON(configPath, &cfg); err != nil {
-		return fmt.Errorf("config: %w", err)
-	}
-	m.Config = &cfg
-	m.Layers = make([]*Block, cfg.NumHiddenLayers)
-
-	// Load weights from tensor blobs
-	weights, err := imagegen.LoadWeightsFromManifest(manifest, "text_encoder")
-	if err != nil {
-		return fmt.Errorf("weights: %w", err)
-	}
-	if err := weights.Load(0); err != nil {
-		return fmt.Errorf("load weights: %w", err)
-	}
-	defer weights.ReleaseAll()
-
-	return m.loadWeights(weights)
-}
-
-// loadWeights loads weights from any WeightSource into the model
-func (m *TextEncoder) loadWeights(weights safetensors.WeightSource) error {
-	if err := safetensors.LoadModule(m, weights, ""); err != nil {
-		return fmt.Errorf("load module: %w", err)
-	}
-	m.initComputedFields()
-	fmt.Println("✓")
-	return nil
-}
-
-// initComputedFields initializes computed fields after loading weights
-func (m *TextEncoder) initComputedFields() {
-	cfg := m.Config
-	m.FinalNorm.Eps = cfg.RMSNormEps
-	for _, block := range m.Layers {
-		// Attention
-		block.Attention.NHeads = cfg.NumAttentionHeads
-		block.Attention.NKVHeads = cfg.NumKeyValueHeads
-		block.Attention.HeadDim = cfg.HeadDim
-		block.Attention.Scale = float32(1.0 / math.Sqrt(float64(cfg.HeadDim)))
-		block.Attention.RopeTheta = cfg.RopeTheta
-		block.Attention.QNorm.Eps = cfg.RMSNormEps
-		block.Attention.KNorm.Eps = cfg.RMSNormEps
-		// Block norms
-		block.InputLayerNorm.Eps = cfg.RMSNormEps
-		block.PostAttnLayerNorm.Eps = cfg.RMSNormEps
-	}
-}
-
-// Forward encodes text tokens with provided attention mask (LxL) and mask mode.
-func (te *TextEncoder) Forward(tokens *mlx.Array, attnMask *mlx.Array, maskMode string) *mlx.Array {
-	h := te.EmbedTokens.Forward(tokens)
-	eps := te.RMSNormEps
-
-	for _, layer := range te.Layers {
-		h = layer.Forward(h, eps, attnMask, maskMode)
-	}
-
-	// Apply final RMS norm
-	h = te.FinalNorm.Forward(h, eps)
-
-	return h
-}
-
-// ForwardWithLayerOutputs encodes text tokens and returns hidden states from specified layers.
-// This is used by Flux2 which needs embeddings from specific intermediate layers.
-func (te *TextEncoder) ForwardWithLayerOutputs(tokens *mlx.Array, layerIndices []int, attnMask *mlx.Array, maskMode string) []*mlx.Array {
-	h := te.EmbedTokens.Forward(tokens)
-	eps := te.RMSNormEps
-
-	outputs := make([]*mlx.Array, len(layerIndices))
-	layerSet := make(map[int]int)
-	for i, idx := range layerIndices {
-		layerSet[idx] = i
-	}
-
-	for i, layer := range te.Layers {
-		h = layer.Forward(h, eps, attnMask, maskMode)
-		if outIdx, ok := layerSet[i]; ok {
-			outputs[outIdx] = h
-		}
-	}
-
-	return outputs
-}
-
-// ApplyChatTemplate wraps prompt in Qwen3 chat format.
-// If think is true, adds the <think></think> block after the assistant tag
-// (matches tokenizer.apply_chat_template with enable_thinking=False in Python).
-func ApplyChatTemplate(prompt string, think bool) string {
-	base := "<|im_start|>user\n" + prompt + "<|im_end|>\n<|im_start|>assistant\n"
-	if think {
-		return base + "<think>\n\n</think>\n\n"
-	}
-	return base
-}
-
-// EncodePrompt encodes a text prompt using the tokenizer and encoder.
-// If think is true, includes the <think></think> block in the chat template.
-func (te *TextEncoder) EncodePrompt(tok *tokenizer.Tokenizer, prompt string, maxLen int, think bool) (*mlx.Array, *mlx.Array) {
-	formattedPrompt := ApplyChatTemplate(prompt, think)
-
-	tokens := tok.Encode(formattedPrompt, false)
-
-	if len(tokens) > maxLen {
-		tokens = tokens[:maxLen]
-	}
-
-	maskData := make([]float32, maxLen)
-	for i := 0; i < len(tokens); i++ {
-		maskData[i] = 1.0
-	}
-
-	// Get PAD token (different from EOS for Qwen3)
-	padToken := tok.PAD()
-	if padToken < 0 {
-		padToken = tok.EOS() // fallback
-	}
-
-	paddedTokens := make([]int32, maxLen)
-	copy(paddedTokens, tokens)
-	for i := len(tokens); i < maxLen; i++ {
-		paddedTokens[i] = padToken
-	}
-
-	tokensArr := mlx.NewArrayInt32(paddedTokens, []int32{1, int32(maxLen)})
-	maskArr := mlx.NewArray(maskData, []int32{1, int32(maxLen)})
-
-	// Build combined causal + PAD mask [L, L]
-	// mask[i,j] = 0 if (j <= i AND valid[j]) else -inf
-	L := int32(maxLen)
-	validLen := int32(len(tokens))
-	combinedMaskData := make([]float32, L*L)
-	negInf := float32(-1e9)
-	for i := int32(0); i < L; i++ {
-		for j := int32(0); j < L; j++ {
-			idx := i*L + j
-			if j <= i && j < validLen {
-				combinedMaskData[idx] = 0
-			} else {
-				combinedMaskData[idx] = negInf
-			}
-		}
-	}
-	maskMat := mlx.NewArray(combinedMaskData, []int32{L, L})
-
-	embeddings := te.Forward(tokensArr, maskMat, "")
-
-	return embeddings, maskArr
-}
-
-// EncodePromptWithLayers encodes a text prompt and returns embeddings from specified layers.
-// Used by Flux2 which concatenates embeddings from multiple intermediate layers.
-// If think is true, includes the <think></think> block in the chat template.
-// Returns embeddings and padded sequence length.
-func (te *TextEncoder) EncodePromptWithLayers(tok *tokenizer.Tokenizer, prompt string, maxLen int, layerIndices []int, think bool) (*mlx.Array, int32) {
-	formattedPrompt := ApplyChatTemplate(prompt, think)
-	tokens := tok.Encode(formattedPrompt, false)
-
-	if len(tokens) > maxLen {
-		tokens = tokens[:maxLen]
-	}
-
-	// Pad to maxLen
-	padToken := tok.PAD()
-	if padToken < 0 {
-		padToken = tok.EOS() // fallback
-	}
-	padded := make([]int32, maxLen)
-	copy(padded, tokens)
-	for i := len(tokens); i < maxLen; i++ {
-		padded[i] = padToken
-	}
-	tokensArr := mlx.NewArrayInt32(padded, []int32{1, int32(maxLen)})
-
-	// Build combined causal + PAD mask [L, L]
-	// mask[i,j] = 0 if (j <= i AND valid[j]) else -inf
-	// This combines causal masking with PAD token masking
-	L := int32(maxLen)
-	validLen := int32(len(tokens))
-	maskData := make([]float32, L*L)
-	negInf := float32(-1e9)
-	for i := int32(0); i < L; i++ {
-		for j := int32(0); j < L; j++ {
-			idx := i*L + j
-			if j <= i && j < validLen {
-				maskData[idx] = 0 // allowed: causal OK and not PAD
-			} else {
-				maskData[idx] = negInf // blocked: future or PAD
-			}
-		}
-	}
-	maskMat := mlx.NewArray(maskData, []int32{L, L})
-
-	layerOutputs := te.ForwardWithLayerOutputs(tokensArr, layerIndices, maskMat, "")
-
-	// Concatenate layer outputs along the hidden dimension
-	// Each output is [B, L, hidden_dim], result is [B, L, num_layers * hidden_dim]
-	embeddings := mlx.Concatenate(layerOutputs, 2)
-
-	// Return embeddings and padded length
-	return embeddings, int32(maxLen)
-}

x/imagegen/models/qwen_image/pipeline_test.go

@@ -3,33 +3,12 @@
 package qwen_image
 
 import (
-	"fmt"
 	"os"
-	"path/filepath"
-	"runtime"
 	"testing"
 
 	"github.com/ollama/ollama/x/imagegen/mlx"
 )
 
-// TestMain initializes MLX before running tests.
-// If MLX libraries are not available, tests are skipped.
-func TestMain(m *testing.M) {
-	// Change to repo root so ./build/lib/ollama/ path works
-	_, thisFile, _, _ := runtime.Caller(0)
-	repoRoot := filepath.Join(filepath.Dir(thisFile), "..", "..", "..", "..")
-	if err := os.Chdir(repoRoot); err != nil {
-		fmt.Printf("Failed to change to repo root: %v\n", err)
-		os.Exit(1)
-	}
-
-	if err := mlx.InitMLX(); err != nil {
-		fmt.Printf("Skipping qwen_image tests: %v\n", err)
-		os.Exit(0)
-	}
-	os.Exit(m.Run())
-}
-
 // TestPipelineOutput runs the full pipeline (integration test).
 // Skips if model weights not found. Requires ~50GB VRAM.
 func TestPipelineOutput(t *testing.T) {

x/imagegen/models/qwen_image/qwen_image.go

@@ -17,13 +17,13 @@ import (
 // GenerateConfig holds all options for image generation.
 type GenerateConfig struct {
 	Prompt         string
-	NegativePrompt string                // Empty = no CFG
-	CFGScale       float32               // Only used if NegativePrompt is set (default: 4.0)
-	Width          int32                 // Image width (default: 1024)
-	Height         int32                 // Image height (default: 1024)
-	Steps          int                   // Denoising steps (default: 30)
-	Seed           int64                 // Random seed
-	Progress       func(step, totalSteps int) // Optional progress callback
+	NegativePrompt string       // Empty = no CFG
+	CFGScale       float32      // Only used if NegativePrompt is set (default: 4.0)
+	Width          int32        // Image width (default: 1024)
+	Height         int32        // Image height (default: 1024)
+	Steps          int          // Denoising steps (default: 30)
+	Seed           int64        // Random seed
+	Progress       ProgressFunc // Optional progress callback
 
 	// Layer caching (DeepCache/Learning-to-Cache speedup)
 	LayerCache    bool // Enable layer caching (default: false)
@@ -31,6 +31,9 @@ type GenerateConfig struct {
 	CacheLayers   int  // Number of shallow layers to cache (default: 25)
 }
 
+// ProgressFunc is called during generation with step progress.
+type ProgressFunc func(step, totalSteps int)
+
 // Model represents a Qwen-Image diffusion model.
 type Model struct {
 	ModelPath   string
@@ -114,7 +117,7 @@ func (m *Model) Generate(prompt string, width, height int32, steps int, seed int
 }
 
 // GenerateWithProgress creates an image with progress callback.
-func (m *Model) GenerateWithProgress(prompt string, width, height int32, steps int, seed int64, progress func(step, totalSteps int)) (*mlx.Array, error) {
+func (m *Model) GenerateWithProgress(prompt string, width, height int32, steps int, seed int64, progress ProgressFunc) (*mlx.Array, error) {
 	return m.GenerateFromConfig(&GenerateConfig{
 		Prompt:   prompt,
 		Width:    width,
@@ -126,7 +129,7 @@ func (m *Model) GenerateWithProgress(prompt string, width, height int32, steps i
 }
 
 // GenerateWithCFG creates an image with classifier-free guidance.
-func (m *Model) GenerateWithCFG(prompt, negativePrompt string, width, height int32, steps int, seed int64, cfgScale float32, progress func(step, totalSteps int)) (*mlx.Array, error) {
+func (m *Model) GenerateWithCFG(prompt, negativePrompt string, width, height int32, steps int, seed int64, cfgScale float32, progress ProgressFunc) (*mlx.Array, error) {
 	return m.GenerateFromConfig(&GenerateConfig{
 		Prompt:         prompt,
 		NegativePrompt: negativePrompt,
@@ -169,7 +172,7 @@ func (m *Model) generate(cfg *GenerateConfig) (*mlx.Array, error) {
 		cfg.Height = 1024
 	}
 	if cfg.Steps <= 0 {
-		cfg.Steps = 50
+		cfg.Steps = 30
 	}
 	if cfg.CFGScale <= 0 {
 		cfg.CFGScale = 4.0

x/imagegen/models/qwen_image_edit/qwen_image_edit.go

@@ -18,15 +18,18 @@ import (
 // GenerateConfig holds all options for image editing.
 type GenerateConfig struct {
 	Prompt         string
-	NegativePrompt string                // Unconditional prompt for CFG (empty string "" is valid)
-	CFGScale       float32               // CFG enabled when > 1.0 (default: 4.0)
-	Width          int32                 // Output width (default: from input image)
-	Height         int32                 // Output height (default: from input image)
-	Steps          int                   // Denoising steps (default: 50)
-	Seed           int64                 // Random seed
-	Progress       func(step, totalSteps int) // Optional progress callback
+	NegativePrompt string       // Unconditional prompt for CFG (empty string "" is valid)
+	CFGScale       float32      // CFG enabled when > 1.0 (default: 4.0)
+	Width          int32        // Output width (default: from input image)
+	Height         int32        // Output height (default: from input image)
+	Steps          int          // Denoising steps (default: 50)
+	Seed           int64        // Random seed
+	Progress       ProgressFunc // Optional progress callback
 }
 
+// ProgressFunc is called during generation with step progress.
+type ProgressFunc func(step, totalSteps int)
+
 // Model represents a Qwen-Image-Edit diffusion model.
 type Model struct {
 	ModelPath     string

x/imagegen/models/qwen_image_edit/rope_test.go

@@ -3,35 +3,13 @@
 package qwen_image_edit
 
 import (
-	"fmt"
 	"math"
-	"os"
-	"path/filepath"
-	"runtime"
 	"testing"
 
 	"github.com/ollama/ollama/x/imagegen/mlx"
 	"github.com/ollama/ollama/x/imagegen/models/qwen_image"
 )
 
-// TestMain initializes MLX before running tests.
-// If MLX libraries are not available, tests are skipped.
-func TestMain(m *testing.M) {
-	// Change to repo root so ./build/lib/ollama/ path works
-	_, thisFile, _, _ := runtime.Caller(0)
-	repoRoot := filepath.Join(filepath.Dir(thisFile), "..", "..", "..", "..")
-	if err := os.Chdir(repoRoot); err != nil {
-		fmt.Printf("Failed to change to repo root: %v\n", err)
-		os.Exit(1)
-	}
-
-	if err := mlx.InitMLX(); err != nil {
-		fmt.Printf("Skipping qwen_image_edit tests: %v\n", err)
-		os.Exit(0)
-	}
-	os.Exit(m.Run())
-}
-
 // TestComputeAxisFreqs verifies frequency computation matches Python reference
 func TestComputeAxisFreqs(t *testing.T) {
 	theta := float64(10000)

x/imagegen/models/zimage/text_encoder.go

@@ -3,17 +3,287 @@
 package zimage
 
 import (
-	"github.com/ollama/ollama/x/imagegen/models/qwen3"
+	"fmt"
+	"math"
+
+	"github.com/ollama/ollama/x/imagegen"
+	"github.com/ollama/ollama/x/imagegen/mlx"
+	"github.com/ollama/ollama/x/imagegen/nn"
+	"github.com/ollama/ollama/x/imagegen/safetensors"
+	"github.com/ollama/ollama/x/imagegen/tokenizer"
 )
 
-// Re-export types from shared qwen3 package for backwards compatibility
-type (
-	Qwen3Config      = qwen3.Config
-	Qwen3Attention   = qwen3.Attention
-	Qwen3MLP         = qwen3.MLP
-	Qwen3Block       = qwen3.Block
-	Qwen3TextEncoder = qwen3.TextEncoder
-)
+// Qwen3Config holds Qwen3 text encoder configuration
+type Qwen3Config struct {
+	HiddenSize        int32   `json:"hidden_size"`
+	NumHiddenLayers   int32   `json:"num_hidden_layers"`
+	IntermediateSize  int32   `json:"intermediate_size"`
+	NumAttentionHeads int32   `json:"num_attention_heads"`
+	NumKeyValueHeads  int32   `json:"num_key_value_heads"`
+	VocabSize         int32   `json:"vocab_size"`
+	RMSNormEps        float32 `json:"rms_norm_eps"`
+	RopeTheta         float32 `json:"rope_theta"`
+	HeadDim           int32   `json:"head_dim"`
+}
+
+// Qwen3Attention implements Qwen3 attention with QK norms
+type Qwen3Attention struct {
+	QProj nn.LinearLayer `weight:"q_proj"`
+	KProj nn.LinearLayer `weight:"k_proj"`
+	VProj nn.LinearLayer `weight:"v_proj"`
+	OProj nn.LinearLayer `weight:"o_proj"`
+	QNorm *nn.RMSNorm    `weight:"q_norm"`
+	KNorm *nn.RMSNorm    `weight:"k_norm"`
+	// Computed fields
+	NHeads    int32
+	NKVHeads  int32
+	HeadDim   int32
+	Scale     float32
+	RopeTheta float32
+}
+
+// applyRoPEQwen3 applies the custom RoPE for Qwen3 text encoder
+func applyRoPEQwen3(x *mlx.Array, seqLen int32, theta float32) *mlx.Array {
+	shape := x.Shape()
+	B := shape[0]
+	L := shape[1]
+	H := shape[2]
+	D := shape[3]
+	half := D / 2
+
+	freqsArr := make([]float32, half)
+	logTheta := float32(math.Log(float64(theta)))
+	for i := int32(0); i < half; i++ {
+		freqsArr[i] = float32(math.Exp(float64(-logTheta * float32(i) / float32(half))))
+	}
+	freqs := mlx.NewArray(freqsArr, []int32{half})
+
+	posArr := make([]float32, seqLen)
+	for i := int32(0); i < seqLen; i++ {
+		posArr[i] = float32(i)
+	}
+	pos := mlx.NewArray(posArr, []int32{seqLen})
+
+	posExpanded := mlx.Reshape(pos, seqLen, 1)
+	freqsExpanded := mlx.Reshape(freqs, 1, half)
+	args := mlx.Mul(posExpanded, freqsExpanded)
+
+	cosVals := mlx.Cos(args)
+	sinVals := mlx.Sin(args)
+	cosVals = mlx.Reshape(cosVals, seqLen, 1, half)
+	sinVals = mlx.Reshape(sinVals, seqLen, 1, half)
+
+	x1 := mlx.Slice(x, []int32{0, 0, 0, 0}, []int32{B, L, H, half})
+	x2 := mlx.Slice(x, []int32{0, 0, 0, half}, []int32{B, L, H, D})
+
+	part1 := mlx.Sub(mlx.Mul(x1, cosVals), mlx.Mul(x2, sinVals))
+	part2 := mlx.Add(mlx.Mul(x1, sinVals), mlx.Mul(x2, cosVals))
+
+	return mlx.Concatenate([]*mlx.Array{part1, part2}, 3)
+}
+
+// Forward computes attention with causal masking
+func (attn *Qwen3Attention) Forward(x *mlx.Array) *mlx.Array {
+	shape := x.Shape()
+	B := shape[0]
+	L := shape[1]
+
+	q := attn.QProj.Forward(x)
+	k := attn.KProj.Forward(x)
+	v := attn.VProj.Forward(x)
+
+	q = mlx.Reshape(q, B, L, attn.NHeads, attn.HeadDim)
+	k = mlx.Reshape(k, B, L, attn.NKVHeads, attn.HeadDim)
+	v = mlx.Reshape(v, B, L, attn.NKVHeads, attn.HeadDim)
+
+	// QK norm uses 1e-6 hardcoded (Qwen3 specific)
+	q = attn.QNorm.Forward(q, 1e-6)
+	k = attn.KNorm.Forward(k, 1e-6)
+
+	q = applyRoPEQwen3(q, L, attn.RopeTheta)
+	k = applyRoPEQwen3(k, L, attn.RopeTheta)
+
+	q = mlx.Transpose(q, 0, 2, 1, 3)
+	k = mlx.Transpose(k, 0, 2, 1, 3)
+	v = mlx.Transpose(v, 0, 2, 1, 3)
+
+	if attn.NKVHeads < attn.NHeads {
+		repeats := attn.NHeads / attn.NKVHeads
+		k = repeatKV(k, repeats)
+		v = repeatKV(v, repeats)
+	}
+
+	out := mlx.ScaledDotProductAttention(q, k, v, attn.Scale, true)
+
+	out = mlx.Transpose(out, 0, 2, 1, 3)
+	out = mlx.Reshape(out, B, L, attn.NHeads*attn.HeadDim)
+
+	out = attn.OProj.Forward(out)
+
+	return out
+}
+
+// repeatKV repeats key/value heads for GQA
+func repeatKV(x *mlx.Array, repeats int32) *mlx.Array {
+	if repeats == 1 {
+		return x
+	}
+	shape := x.Shape()
+	x = mlx.ExpandDims(x, 2)
+	x = mlx.Tile(x, []int32{1, 1, repeats, 1, 1})
+	return mlx.Reshape(x, shape[0], shape[1]*repeats, shape[2], shape[3])
+}
+
+// Qwen3MLP implements Qwen3 SwiGLU MLP
+type Qwen3MLP struct {
+	GateProj nn.LinearLayer `weight:"gate_proj"`
+	UpProj   nn.LinearLayer `weight:"up_proj"`
+	DownProj nn.LinearLayer `weight:"down_proj"`
+}
+
+// Forward applies the MLP
+func (m *Qwen3MLP) Forward(x *mlx.Array) *mlx.Array {
+	gate := m.GateProj.Forward(x)
+	gate = mlx.SiLU(gate)
+	up := m.UpProj.Forward(x)
+	h := mlx.Mul(gate, up)
+	return m.DownProj.Forward(h)
+}
+
+// Qwen3Block represents a single Qwen3 transformer block
+type Qwen3Block struct {
+	Attention         *Qwen3Attention `weight:"self_attn"`
+	MLP               *Qwen3MLP       `weight:"mlp"`
+	InputLayerNorm    *nn.RMSNorm     `weight:"input_layernorm"`
+	PostAttnLayerNorm *nn.RMSNorm     `weight:"post_attention_layernorm"`
+}
+
+// Forward applies the Qwen3 block
+func (qb *Qwen3Block) Forward(x *mlx.Array, eps float32) *mlx.Array {
+	h := qb.InputLayerNorm.Forward(x, eps)
+	attnOut := qb.Attention.Forward(h)
+	x = mlx.Add(x, attnOut)
+
+	h = qb.PostAttnLayerNorm.Forward(x, eps)
+	mlpOut := qb.MLP.Forward(h)
+	x = mlx.Add(x, mlpOut)
+
+	return x
+}
+
+// Qwen3TextEncoder is the full Qwen3 encoder for Z-Image
+type Qwen3TextEncoder struct {
+	EmbedTokens *nn.Embedding   `weight:"model.embed_tokens"`
+	Layers      []*Qwen3Block   `weight:"model.layers"`
+	FinalNorm   *nn.RMSNorm     `weight:"model.norm"`
+	*Qwen3Config
+}
+
+// Load loads the Qwen3 text encoder from ollama blob storage.
+func (m *Qwen3TextEncoder) Load(manifest *imagegen.ModelManifest) error {
+	fmt.Print("  Loading text encoder... ")
+
+	// Load config from blob
+	var cfg Qwen3Config
+	if err := manifest.ReadConfigJSON("text_encoder/config.json", &cfg); err != nil {
+		return fmt.Errorf("config: %w", err)
+	}
+	m.Qwen3Config = &cfg
+	m.Layers = make([]*Qwen3Block, cfg.NumHiddenLayers)
+
+	// Load weights from tensor blobs
+	weights, err := imagegen.LoadWeightsFromManifest(manifest, "text_encoder")
+	if err != nil {
+		return fmt.Errorf("weights: %w", err)
+	}
+	if err := weights.Load(0); err != nil {
+		return fmt.Errorf("load weights: %w", err)
+	}
+	defer weights.ReleaseAll()
+
+	return m.loadWeights(weights)
+}
+
+// loadWeights loads weights from any WeightSource into the model
+func (m *Qwen3TextEncoder) loadWeights(weights safetensors.WeightSource) error {
+	if err := safetensors.LoadModule(m, weights, ""); err != nil {
+		return fmt.Errorf("load module: %w", err)
+	}
+	m.initComputedFields()
+	fmt.Println("✓")
+	return nil
+}
+
+// initComputedFields initializes computed fields after loading weights
+func (m *Qwen3TextEncoder) initComputedFields() {
+	cfg := m.Qwen3Config
+	m.FinalNorm.Eps = cfg.RMSNormEps
+	for _, block := range m.Layers {
+		// Attention
+		block.Attention.NHeads = cfg.NumAttentionHeads
+		block.Attention.NKVHeads = cfg.NumKeyValueHeads
+		block.Attention.HeadDim = cfg.HeadDim
+		block.Attention.Scale = float32(1.0 / math.Sqrt(float64(cfg.HeadDim)))
+		block.Attention.RopeTheta = cfg.RopeTheta
+		block.Attention.QNorm.Eps = cfg.RMSNormEps
+		block.Attention.KNorm.Eps = cfg.RMSNormEps
+		// Block norms
+		block.InputLayerNorm.Eps = cfg.RMSNormEps
+		block.PostAttnLayerNorm.Eps = cfg.RMSNormEps
+	}
+}
+
+// Forward encodes text tokens
+func (te *Qwen3TextEncoder) Forward(tokens *mlx.Array) *mlx.Array {
+	h := te.EmbedTokens.Forward(tokens)
+	eps := te.RMSNormEps
+
+	for _, layer := range te.Layers {
+		h = layer.Forward(h, eps)
+	}
+
+	// Apply final RMS norm
+	h = te.FinalNorm.Forward(h, eps)
+
+	return h
+}
 
 // ApplyChatTemplate wraps prompt in Qwen3 chat format
-var ApplyChatTemplate = qwen3.ApplyChatTemplate
+func ApplyChatTemplate(prompt string) string {
+	return "<|im_start|>user\n" + prompt + "<|im_end|>\n<|im_start|>assistant\n"
+}
+
+// EncodePrompt encodes a text prompt using the tokenizer and encoder
+func (te *Qwen3TextEncoder) EncodePrompt(tok *tokenizer.Tokenizer, prompt string, maxLen int) (*mlx.Array, *mlx.Array) {
+	formattedPrompt := ApplyChatTemplate(prompt)
+
+	tokens := tok.Encode(formattedPrompt, false)
+
+	if len(tokens) > maxLen {
+		tokens = tokens[:maxLen]
+	}
+
+	maskData := make([]float32, maxLen)
+	for i := 0; i < len(tokens); i++ {
+		maskData[i] = 1.0
+	}
+
+	// Get PAD token (different from EOS for Qwen3)
+	padToken := tok.PAD()
+	if padToken < 0 {
+		padToken = tok.EOS() // fallback
+	}
+
+	paddedTokens := make([]int32, maxLen)
+	copy(paddedTokens, tokens)
+	for i := len(tokens); i < maxLen; i++ {
+		paddedTokens[i] = padToken
+	}
+
+	tokensArr := mlx.NewArrayInt32(paddedTokens, []int32{1, int32(maxLen)})
+	maskArr := mlx.NewArray(maskData, []int32{1, int32(maxLen)})
+
+	embeddings := te.Forward(tokensArr)
+
+	return embeddings, maskArr
+}

x/imagegen/models/zimage/vae.go

@@ -9,7 +9,6 @@ import (
 	"github.com/ollama/ollama/x/imagegen"
 	"github.com/ollama/ollama/x/imagegen/mlx"
 	"github.com/ollama/ollama/x/imagegen/safetensors"
-	"github.com/ollama/ollama/x/imagegen/vae"
 )
 
 // VAEConfig holds VAE decoder configuration
@@ -637,9 +636,6 @@ type VAEDecoder struct {
 	UpBlocks    []*UpDecoderBlock2D
 	ConvNormOut *GroupNormLayer
 	ConvOut     *Conv2D
-
-	// Tiling configuration (nil = no tiling)
-	Tiling *vae.TilingConfig
 }
 
 // Load loads the VAE decoder from ollama blob storage.
@@ -734,60 +730,45 @@ func (m *VAEDecoder) loadWeights(weights safetensors.WeightSource, cfg *VAEConfi
 
 // Decode decodes latents to images.
 // Input latents are in NCHW format, output is in NCHW format.
-// If Tiling is set, uses tiled decoding to reduce memory for large images.
-func (v *VAEDecoder) Decode(latents *mlx.Array) *mlx.Array {
+// Internally uses NHWC format (MLX native) for all operations.
+func (vae *VAEDecoder) Decode(latents *mlx.Array) *mlx.Array {
 	// Scale latents
-	z := mlx.DivScalar(latents, v.Config.ScalingFactor)
-	z = mlx.AddScalar(z, v.Config.ShiftFactor)
+	z := mlx.DivScalar(latents, vae.Config.ScalingFactor)
+	z = mlx.AddScalar(z, vae.Config.ShiftFactor)
 	// Convert NCHW -> NHWC for internal processing
 	z = mlx.Transpose(z, 0, 2, 3, 1)
-
-	// Use tiled decoding if enabled
-	if v.Tiling != nil {
-		mlx.Eval(z)
-		return vae.DecodeTiled(z, v.Tiling, v.decodeTile)
-	}
-
-	// Direct decode
-	h := v.decodeTile(z)
-	h = mlx.ClipScalar(h, 0.0, 1.0, true, true)
-	// Convert NHWC -> NCHW for output
-	h = mlx.Transpose(h, 0, 3, 1, 2)
-	mlx.Eval(h)
-	return h
-}
-
-// decodeTile decodes a single latent tile to pixels.
-// Input: [B, H, W, C] latent tile in NHWC format (already scaled)
-// Output: [B, H*8, W*8, 3] pixel tile in NHWC format
-func (v *VAEDecoder) decodeTile(z *mlx.Array) *mlx.Array {
-	h := v.ConvIn.Forward(z)
+	h := vae.ConvIn.Forward(z)
 	mlx.Eval(h)
 
 	prev := h
-	h = v.MidBlock.Forward(h)
+	h = vae.MidBlock.Forward(h)
 	prev.Free()
 
-	for _, upBlock := range v.UpBlocks {
+	for _, upBlock := range vae.UpBlocks {
 		prev = h
 		h = upBlock.Forward(h)
 		prev.Free()
 	}
 
 	prev = h
-	h = v.ConvNormOut.Forward(h)
+	h = vae.ConvNormOut.Forward(h)
 	mlx.Eval(h) // Eval after GroupNorm to avoid grid dimension issues
 	prev.Free()
 
 	prev = h
 	h = mlx.SiLU(h)
-	h = v.ConvOut.Forward(h)
+	h = vae.ConvOut.Forward(h)
 	mlx.Eval(h)
 	prev.Free()
 
 	// VAE outputs [-1, 1], convert to [0, 1]
 	h = mlx.MulScalar(h, 0.5)
 	h = mlx.AddScalar(h, 0.5)
+	h = mlx.ClipScalar(h, 0.0, 1.0, true, true)
+
+	// Convert NHWC -> NCHW for output
+	h = mlx.Transpose(h, 0, 3, 1, 2)
+	mlx.Eval(h)
 
 	return h
 }

x/imagegen/models/zimage/zimage.go

@@ -12,20 +12,19 @@ import (
 	"github.com/ollama/ollama/x/imagegen/cache"
 	"github.com/ollama/ollama/x/imagegen/mlx"
 	"github.com/ollama/ollama/x/imagegen/tokenizer"
-	"github.com/ollama/ollama/x/imagegen/vae"
 )
 
 // GenerateConfig holds all options for image generation.
 type GenerateConfig struct {
 	Prompt         string
-	NegativePrompt string                // Empty = no CFG
-	CFGScale       float32               // Only used if NegativePrompt is set (default: 4.0)
-	Width          int32                 // Image width (default: 1024)
-	Height         int32                 // Image height (default: 1024)
-	Steps          int                   // Denoising steps (default: 9 for turbo)
-	Seed           int64                 // Random seed
-	Progress       func(step, totalSteps int) // Optional progress callback
-	CapturePath    string                // GPU capture path (debug)
+	NegativePrompt string       // Empty = no CFG
+	CFGScale       float32      // Only used if NegativePrompt is set (default: 4.0)
+	Width          int32        // Image width (default: 1024)
+	Height         int32        // Image height (default: 1024)
+	Steps          int          // Denoising steps (default: 9 for turbo)
+	Seed           int64        // Random seed
+	Progress       ProgressFunc // Optional progress callback
+	CapturePath    string       // GPU capture path (debug)
 
 	// TeaCache options (timestep embedding aware caching)
 	TeaCache          bool    // TeaCache is always enabled for faster inference
@@ -35,6 +34,9 @@ type GenerateConfig struct {
 	FusedQKV bool // Enable fused QKV projection (default: false)
 }
 
+// ProgressFunc is called during generation with step progress.
+type ProgressFunc func(step, totalSteps int)
+
 // Model represents a Z-Image diffusion model.
 type Model struct {
 	ModelName   string
@@ -91,7 +93,7 @@ func (m *Model) Load(modelName string) error {
 
 	// Load text encoder
 	m.TextEncoder = &Qwen3TextEncoder{}
-	if err := m.TextEncoder.Load(manifest, "text_encoder/config.json"); err != nil {
+	if err := m.TextEncoder.Load(manifest); err != nil {
 		return fmt.Errorf("text encoder: %w", err)
 	}
 	mlx.Eval(mlx.Collect(m.TextEncoder)...)
@@ -137,7 +139,7 @@ func (m *Model) Generate(prompt string, width, height int32, steps int, seed int
 }
 
 // GenerateWithProgress creates an image with progress callback.
-func (m *Model) GenerateWithProgress(prompt string, width, height int32, steps int, seed int64, progress func(step, totalSteps int)) (*mlx.Array, error) {
+func (m *Model) GenerateWithProgress(prompt string, width, height int32, steps int, seed int64, progress ProgressFunc) (*mlx.Array, error) {
 	return m.GenerateFromConfig(context.Background(), &GenerateConfig{
 		Prompt:   prompt,
 		Width:    width,
@@ -149,7 +151,7 @@ func (m *Model) GenerateWithProgress(prompt string, width, height int32, steps i
 }
 
 // GenerateWithCFG creates an image with classifier-free guidance.
-func (m *Model) GenerateWithCFG(prompt, negativePrompt string, width, height int32, steps int, seed int64, cfgScale float32, progress func(step, totalSteps int)) (*mlx.Array, error) {
+func (m *Model) GenerateWithCFG(prompt, negativePrompt string, width, height int32, steps int, seed int64, cfgScale float32, progress ProgressFunc) (*mlx.Array, error) {
 	return m.GenerateFromConfig(context.Background(), &GenerateConfig{
 		Prompt:         prompt,
 		NegativePrompt: negativePrompt,
@@ -177,16 +179,9 @@ func (m *Model) GenerateFromConfig(ctx context.Context, cfg *GenerateConfig) (*m
 	return result, nil
 }
 
-// GenerateImage implements runner.ImageModel interface.
-func (m *Model) GenerateImage(ctx context.Context, prompt string, width, height int32, steps int, seed int64, progress func(step, total int)) (*mlx.Array, error) {
-	return m.GenerateFromConfig(ctx, &GenerateConfig{
-		Prompt:   prompt,
-		Width:    width,
-		Height:   height,
-		Steps:    steps,
-		Seed:     seed,
-		Progress: progress,
-	})
+// GenerateImage implements model.ImageModel interface.
+func (m *Model) GenerateImage(ctx context.Context, prompt string, width, height int32, steps int, seed int64) (*mlx.Array, error) {
+	return m.Generate(prompt, width, height, steps, seed)
 }
 
 // generate is the internal denoising pipeline.
@@ -199,7 +194,7 @@ func (m *Model) generate(ctx context.Context, cfg *GenerateConfig) (*mlx.Array,
 		cfg.Height = 1024
 	}
 	if cfg.Steps <= 0 {
-		cfg.Steps = 9 // Z-Image turbo default
+		cfg.Steps = 9 // Turbo default
 	}
 	if cfg.CFGScale <= 0 {
 		cfg.CFGScale = 4.0
@@ -227,9 +222,9 @@ func (m *Model) generate(ctx context.Context, cfg *GenerateConfig) (*mlx.Array,
 	// Text encoding with padding to multiple of 32
 	var posEmb, negEmb *mlx.Array
 	{
-		posEmb, _ = m.TextEncoder.EncodePrompt(m.Tokenizer, cfg.Prompt, 512, false)
+		posEmb, _ = m.TextEncoder.EncodePrompt(m.Tokenizer, cfg.Prompt, 512)
 		if useCFG {
-			negEmb, _ = m.TextEncoder.EncodePrompt(m.Tokenizer, cfg.NegativePrompt, 512, false)
+			negEmb, _ = m.TextEncoder.EncodePrompt(m.Tokenizer, cfg.NegativePrompt, 512)
 		}
 
 		// Pad both to same length (multiple of 32)
@@ -453,11 +448,7 @@ func (m *Model) generate(ctx context.Context, cfg *GenerateConfig) (*mlx.Array,
 		teaCache.Free()
 	}
 
-	// VAE decode - enable tiling for larger images to reduce memory
-	// VAE attention is O(n²) on latent pixels, tiling helps significantly
-	if latentH > 64 || latentW > 64 {
-		m.VAEDecoder.Tiling = vae.DefaultTilingConfig()
-	}
+	// VAE decode
 	decoded := m.VAEDecoder.Decode(latents)
 	latents.Free()
 

x/imagegen/nn/nn_test.go

@@ -3,34 +3,12 @@
 package nn
 
 import (
-	"fmt"
 	"math"
-	"os"
-	"path/filepath"
-	"runtime"
 	"testing"
 
 	"github.com/ollama/ollama/x/imagegen/mlx"
 )
 
-// TestMain initializes MLX before running tests.
-// If MLX libraries are not available, tests are skipped.
-func TestMain(m *testing.M) {
-	// Change to repo root so ./build/lib/ollama/ path works
-	_, thisFile, _, _ := runtime.Caller(0)
-	repoRoot := filepath.Join(filepath.Dir(thisFile), "..", "..", "..")
-	if err := os.Chdir(repoRoot); err != nil {
-		fmt.Printf("Failed to change to repo root: %v\n", err)
-		os.Exit(1)
-	}
-
-	if err := mlx.InitMLX(); err != nil {
-		fmt.Printf("Skipping nn tests: %v\n", err)
-		os.Exit(0)
-	}
-	os.Exit(m.Run())
-}
-
 // TestLinearNoBias verifies Linear without bias computes x @ w.T correctly.
 func TestLinearNoBias(t *testing.T) {
 	// Weight: [out=2, in=3] -> transposed at forward time

x/imagegen/quantize.go

@@ -0,0 +1,22 @@
+package imagegen
+
+import (
+	"io"
+	"strings"
+)
+
+// QuantizingTensorLayerCreator creates tensor layers with optional quantization.
+// When quantize is true, returns multiple layers (weight + scales + biases).
+type QuantizingTensorLayerCreator func(r io.Reader, name, dtype string, shape []int32, quantize bool) ([]LayerInfo, error)
+
+// ShouldQuantize returns true if a tensor should be quantized.
+// Quantizes linear weights only, skipping VAE, embeddings, norms, and biases.
+func ShouldQuantize(name, component string) bool {
+	if component == "vae" {
+		return false
+	}
+	if strings.Contains(name, "embed") || strings.Contains(name, "norm") {
+		return false
+	}
+	return strings.HasSuffix(name, ".weight")
+}

x/imagegen/runner/runner.go

@@ -19,7 +19,6 @@ import (
 
 	"github.com/ollama/ollama/x/imagegen"
 	"github.com/ollama/ollama/x/imagegen/mlx"
-	"github.com/ollama/ollama/x/imagegen/models/flux2"
 	"github.com/ollama/ollama/x/imagegen/models/zimage"
 )
 
@@ -37,19 +36,12 @@ type Response struct {
 	Content string `json:"content,omitempty"`
 	Image   string `json:"image,omitempty"` // Base64-encoded PNG
 	Done    bool   `json:"done"`
-	Step    int    `json:"step,omitempty"`
-	Total   int    `json:"total,omitempty"`
-}
-
-// ImageModel is the interface for image generation models
-type ImageModel interface {
-	GenerateImage(ctx context.Context, prompt string, width, height int32, steps int, seed int64, progress func(step, total int)) (*mlx.Array, error)
 }
 
 // Server holds the model and handles requests
 type Server struct {
 	mu        sync.Mutex
-	model     ImageModel
+	model     *zimage.Model
 	modelName string
 }
 
@@ -70,12 +62,6 @@ func Execute(args []string) error {
 		return fmt.Errorf("--port is required")
 	}
 
-	err := mlx.InitMLX()
-	if err != nil {
-		slog.Error("unable to initialize MLX", "error", err)
-		return err
-	}
-	slog.Info("MLX library initialized")
 	slog.Info("starting image runner", "model", *modelName, "port", *port)
 
 	// Check memory requirements before loading
@@ -86,25 +72,10 @@ func Execute(args []string) error {
 			requiredMemory/(1024*1024*1024), availableMemory/(1024*1024*1024))
 	}
 
-	// Detect model type and load appropriate model
-	modelType := imagegen.DetectModelType(*modelName)
-	slog.Info("detected model type", "type", modelType)
-
-	var model ImageModel
-	switch modelType {
-	case "Flux2KleinPipeline":
-		m := &flux2.Model{}
-		if err := m.Load(*modelName); err != nil {
-			return fmt.Errorf("failed to load model: %w", err)
-		}
-		model = m
-	default:
-		// Default to Z-Image for ZImagePipeline, FluxPipeline, etc.
-		m := &zimage.Model{}
-		if err := m.Load(*modelName); err != nil {
-			return fmt.Errorf("failed to load model: %w", err)
-		}
-		model = m
+	// Load model
+	model := &zimage.Model{}
+	if err := model.Load(*modelName); err != nil {
+		return fmt.Errorf("failed to load model: %w", err)
 	}
 
 	server := &Server{
@@ -165,8 +136,16 @@ func (s *Server) completionHandler(w http.ResponseWriter, r *http.Request) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
-	// Model applies its own defaults for width/height/steps
-	// Only seed needs to be set here if not provided
+	// Apply defaults
+	if req.Width <= 0 {
+		req.Width = 1024
+	}
+	if req.Height <= 0 {
+		req.Height = 1024
+	}
+	if req.Steps <= 0 {
+		req.Steps = 9
+	}
 	if req.Seed <= 0 {
 		req.Seed = time.Now().UnixNano()
 	}
@@ -180,19 +159,25 @@ func (s *Server) completionHandler(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	// Generate image using the common interface
+	// Generate image
 	ctx := r.Context()
-	enc := json.NewEncoder(w)
-
-	// Progress callback streams step updates
-	progress := func(step, total int) {
-		resp := Response{Step: step, Total: total}
-		enc.Encode(resp)
-		w.Write([]byte("\n"))
-		flusher.Flush()
-	}
-
-	img, err := s.model.GenerateImage(ctx, req.Prompt, req.Width, req.Height, req.Steps, req.Seed, progress)
+	img, err := s.model.GenerateFromConfig(ctx, &zimage.GenerateConfig{
+		Prompt: req.Prompt,
+		Width:  req.Width,
+		Height: req.Height,
+		Steps:  req.Steps,
+		Seed:   req.Seed,
+		Progress: func(step, total int) {
+			resp := Response{
+				Content: fmt.Sprintf("\rGenerating: step %d/%d", step, total),
+				Done:    false,
+			}
+			data, _ := json.Marshal(resp)
+			w.Write(data)
+			w.Write([]byte("\n"))
+			flusher.Flush()
+		},
+	})
 
 	if err != nil {
 		// Don't send error for cancellation

x/imagegen/safetensors/loader.go

@@ -17,24 +17,6 @@ type WeightSource interface {
 	GetTensor(name string) (*mlx.Array, error)
 	ListTensors() []string
 	HasTensor(name string) bool
-	Quantization() string // Returns "FP4", "FP8", or ""
-}
-
-// quantizationParams returns groupSize, bits, mode for a quantization type.
-// Returns defaults (32, 8, "affine") for unknown types (backward compatibility).
-func quantizationParams(quantization string) (groupSize, bits int, mode string) {
-	switch strings.ToUpper(quantization) {
-	case "FP4":
-		return 32, 4, "affine"
-	default:
-		return 32, 8, "affine" // FP8 or unknown
-	}
-}
-
-// Transformer allows structs to transform weight arrays before assignment.
-// Implement this to apply operations like transpose during loading.
-type Transformer interface {
-	Transform(field string, arr *mlx.Array) *mlx.Array
 }
 
 // LoadModule loads weights into a struct using reflection and struct tags.
@@ -154,10 +136,6 @@ func loadStruct(v reflect.Value, weights WeightSource, prefix string, errs *[]st
 					}
 					continue
 				}
-				// Transform before assigning if parent implements Transformer
-				if t, ok := v.Addr().Interface().(Transformer); ok {
-					arr = t.Transform(field.Name, arr)
-				}
 				fieldVal.Set(reflect.ValueOf(arr))
 				continue
 			}
@@ -245,21 +223,19 @@ func LoadLinearLayer(weights WeightSource, path string) (nn.LinearLayer, error)
 			qbiases, _ = weights.GetTensor(qbiasPath)
 		}
 
-		groupSize, bits, mode := quantizationParams(weights.Quantization())
-
 		if mlx.MetalIsAvailable() {
 			return &nn.QuantizedLinear{
 				Weight:    weight,
 				Scales:    scales,
 				QBiases:   qbiases,
 				Bias:      bias,
-				GroupSize: groupSize,
-				Bits:      bits,
-				Mode:      mode,
+				GroupSize: 32,
+				Bits:      8,
+				Mode:      "affine",
 			}, nil
 		}
 
-		dequantized := mlx.Dequantize(weight, scales, qbiases, groupSize, bits, mode)
+		dequantized := mlx.Dequantize(weight, scales, qbiases, 32, 8, "affine")
 		return nn.NewLinear(dequantized, bias), nil
 	}
 

x/imagegen/safetensors/safetensors.go

@@ -298,11 +298,6 @@ func (mw *ModelWeights) HasTensor(name string) bool {
 	return ok
 }
 
-// Quantization returns empty string for directory-based weights (not quantized).
-func (mw *ModelWeights) Quantization() string {
-	return ""
-}
-
 // ReleaseAll releases all cached native file handles.
 func (mw *ModelWeights) ReleaseAll() {
 	for path, native := range mw.nativeCache {

x/imagegen/server.go

@@ -25,11 +25,6 @@ import (
 )
 
 // Server wraps an image generation subprocess to implement llm.LlamaServer.
-//
-// This implementation is compatible with Ollama's scheduler and can be loaded/unloaded
-// like any other model. The plan is to eventually bring this into the llm/ package
-// and evolve llm/ to support MLX and multimodal models. For now, keeping the code
-// separate allows for independent iteration on image generation support.
 type Server struct {
 	mu          sync.Mutex
 	cmd         *exec.Cmd
@@ -42,6 +37,22 @@ type Server struct {
 	lastErrLock sync.Mutex
 }
 
+// completionRequest is sent to the subprocess
+type completionRequest struct {
+	Prompt string `json:"prompt"`
+	Width  int32  `json:"width,omitempty"`
+	Height int32  `json:"height,omitempty"`
+	Steps  int    `json:"steps,omitempty"`
+	Seed   int64  `json:"seed,omitempty"`
+}
+
+// completionResponse is received from the subprocess
+type completionResponse struct {
+	Content string `json:"content,omitempty"`
+	Image   string `json:"image,omitempty"`
+	Done    bool   `json:"done"`
+}
+
 // NewServer spawns a new image generation subprocess and waits until it's ready.
 func NewServer(modelName string) (*Server, error) {
 	// Validate platform support before attempting to start
@@ -61,7 +72,7 @@ func NewServer(modelName string) (*Server, error) {
 		port = rand.Intn(65535-49152) + 49152
 	}
 
-	// Get the current executable path (we use the same binary with runner subcommand)
+	// Get the ollama-mlx executable path (in same directory as current executable)
 	exe, err := os.Executable()
 	if err != nil {
 		return nil, fmt.Errorf("unable to lookup executable path: %w", err)
@@ -69,9 +80,10 @@ func NewServer(modelName string) (*Server, error) {
 	if eval, err := filepath.EvalSymlinks(exe); err == nil {
 		exe = eval
 	}
+	mlxExe := filepath.Join(filepath.Dir(exe), "ollama-mlx")
 
-	// Spawn subprocess: ollama runner --image-engine --model <path> --port <port>
-	cmd := exec.Command(exe, "runner", "--image-engine", "--model", modelName, "--port", strconv.Itoa(port))
+	// Spawn subprocess: ollama-mlx runner --image-engine --model <path> --port <port>
+	cmd := exec.Command(mlxExe, "runner", "--image-engine", "--model", modelName, "--port", strconv.Itoa(port))
 	cmd.Env = os.Environ()
 
 	// On Linux, set LD_LIBRARY_PATH to include MLX library directories
@@ -127,13 +139,14 @@ func NewServer(modelName string) (*Server, error) {
 		for scanner.Scan() {
 			line := scanner.Text()
 			slog.Warn("image-runner", "msg", line)
+			// Capture last error line for better error reporting
 			s.lastErrLock.Lock()
 			s.lastErr = line
 			s.lastErrLock.Unlock()
 		}
 	}()
 
-	slog.Info("starting image runner subprocess", "exe", exe, "model", modelName, "port", port)
+	slog.Info("starting ollama-mlx image runner subprocess", "exe", mlxExe, "model", modelName, "port", port)
 	if err := cmd.Start(); err != nil {
 		return nil, fmt.Errorf("failed to start image runner: %w", err)
 	}
@@ -158,6 +171,7 @@ func (s *Server) ModelPath() string {
 	return s.modelName
 }
 
+// Load is called by the scheduler after the server is created.
 func (s *Server) Load(ctx context.Context, systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, requireFull bool) ([]ml.DeviceID, error) {
 	return nil, nil
 }
@@ -190,16 +204,20 @@ func (s *Server) waitUntilRunning() error {
 	for {
 		select {
 		case err := <-s.done:
-			// Include recent stderr lines for better error context
-			errMsg := s.getLastErr()
-			if errMsg != "" {
-				return fmt.Errorf("image runner failed: %s (exit: %v)", errMsg, err)
+			// Include last stderr line for better error context
+			s.lastErrLock.Lock()
+			lastErr := s.lastErr
+			s.lastErrLock.Unlock()
+			if lastErr != "" {
+				return fmt.Errorf("image runner failed: %s (exit: %v)", lastErr, err)
 			}
 			return fmt.Errorf("image runner exited unexpectedly: %w", err)
 		case <-timeout:
-			errMsg := s.getLastErr()
-			if errMsg != "" {
-				return fmt.Errorf("timeout waiting for image runner: %s", errMsg)
+			s.lastErrLock.Lock()
+			lastErr := s.lastErr
+			s.lastErrLock.Unlock()
+			if lastErr != "" {
+				return fmt.Errorf("timeout waiting for image runner: %s", lastErr)
 			}
 			return errors.New("timeout waiting for image runner to start")
 		case <-ticker.C:
@@ -211,41 +229,44 @@ func (s *Server) waitUntilRunning() error {
 	}
 }
 
-// getLastErr returns the last stderr line.
-func (s *Server) getLastErr() string {
-	s.lastErrLock.Lock()
-	defer s.lastErrLock.Unlock()
-	return s.lastErr
+// WaitUntilRunning implements the LlamaServer interface (no-op since NewServer waits).
+func (s *Server) WaitUntilRunning(ctx context.Context) error {
+	return nil
 }
 
-func (s *Server) WaitUntilRunning(ctx context.Context) error { return nil }
-
+// Completion generates an image from the prompt via the subprocess.
 func (s *Server) Completion(ctx context.Context, req llm.CompletionRequest, fn func(llm.CompletionResponse)) error {
-	seed := req.Seed
-	if seed == 0 {
-		seed = time.Now().UnixNano()
-	}
-
-	// Build request for subprocess
-	creq := struct {
-		Prompt string `json:"prompt"`
-		Width  int32  `json:"width,omitempty"`
-		Height int32  `json:"height,omitempty"`
-		Steps  int32  `json:"steps,omitempty"`
-		Seed   int64  `json:"seed,omitempty"`
-	}{
+	// Build request
+	creq := completionRequest{
 		Prompt: req.Prompt,
-		Width:  req.Width,
-		Height: req.Height,
-		Steps:  req.Steps,
-		Seed:   seed,
+		Width:  1024,
+		Height: 1024,
+		Steps:  9,
+		Seed:   time.Now().UnixNano(),
 	}
 
+	if req.Options != nil {
+		if req.Options.NumCtx > 0 && req.Options.NumCtx <= 4096 {
+			creq.Width = int32(req.Options.NumCtx)
+		}
+		if req.Options.NumGPU > 0 && req.Options.NumGPU <= 4096 {
+			creq.Height = int32(req.Options.NumGPU)
+		}
+		if req.Options.NumPredict > 0 && req.Options.NumPredict <= 100 {
+			creq.Steps = req.Options.NumPredict
+		}
+		if req.Options.Seed > 0 {
+			creq.Seed = int64(req.Options.Seed)
+		}
+	}
+
+	// Encode request body
 	body, err := json.Marshal(creq)
 	if err != nil {
 		return err
 	}
 
+	// Send request to subprocess
 	url := fmt.Sprintf("http://127.0.0.1:%d/completion", s.port)
 	httpReq, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body))
 	if err != nil {
@@ -260,36 +281,30 @@ func (s *Server) Completion(ctx context.Context, req llm.CompletionRequest, fn f
 	defer resp.Body.Close()
 
 	if resp.StatusCode != http.StatusOK {
-		return fmt.Errorf("request failed: %d", resp.StatusCode)
+		return fmt.Errorf("completion request failed: %d", resp.StatusCode)
 	}
 
+	// Stream responses - use large buffer for base64 image data
 	scanner := bufio.NewScanner(resp.Body)
 	scanner.Buffer(make([]byte, 1024*1024), 16*1024*1024) // 16MB max
 	for scanner.Scan() {
-		// Parse subprocess response (has singular "image" field)
-		var raw struct {
-			Image   string `json:"image,omitempty"`
-			Content string `json:"content,omitempty"`
-			Done    bool   `json:"done"`
-			Step    int    `json:"step,omitempty"`
-			Total   int    `json:"total,omitempty"`
-		}
-		if err := json.Unmarshal(scanner.Bytes(), &raw); err != nil {
+		var cresp completionResponse
+		if err := json.Unmarshal(scanner.Bytes(), &cresp); err != nil {
 			continue
 		}
 
-		// Convert to llm.CompletionResponse
-		cresp := llm.CompletionResponse{
-			Content:    raw.Content,
-			Done:       raw.Done,
-			Step:       raw.Step,
-			TotalSteps: raw.Total,
-			Image:      raw.Image,
+		content := cresp.Content
+		// If this is the final response with an image, encode it in the content
+		if cresp.Done && cresp.Image != "" {
+			content = "IMAGE_BASE64:" + cresp.Image
 		}
 
-		fn(cresp)
+		fn(llm.CompletionResponse{
+			Content: content,
+			Done:    cresp.Done,
+		})
 		if cresp.Done {
-			return nil
+			break
 		}
 	}
 
@@ -331,18 +346,22 @@ func (s *Server) VRAMByGPU(id ml.DeviceID) uint64 {
 	return s.vramSize
 }
 
+// Embedding is not supported for image generation models.
 func (s *Server) Embedding(ctx context.Context, input string) ([]float32, int, error) {
-	return nil, 0, errors.New("not supported")
+	return nil, 0, errors.New("embedding not supported for image generation models")
 }
 
+// Tokenize is not supported for image generation models.
 func (s *Server) Tokenize(ctx context.Context, content string) ([]int, error) {
-	return nil, errors.New("not supported")
+	return nil, errors.New("tokenize not supported for image generation models")
 }
 
+// Detokenize is not supported for image generation models.
 func (s *Server) Detokenize(ctx context.Context, tokens []int) (string, error) {
-	return "", errors.New("not supported")
+	return "", errors.New("detokenize not supported for image generation models")
 }
 
+// Pid returns the subprocess PID.
 func (s *Server) Pid() int {
 	s.mu.Lock()
 	defer s.mu.Unlock()
@@ -352,9 +371,17 @@ func (s *Server) Pid() int {
 	return -1
 }
 
-func (s *Server) GetPort() int                                       { return s.port }
-func (s *Server) GetDeviceInfos(ctx context.Context) []ml.DeviceInfo { return nil }
+// GetPort returns the subprocess port.
+func (s *Server) GetPort() int {
+	return s.port
+}
 
+// GetDeviceInfos returns nil since we don't track GPU info.
+func (s *Server) GetDeviceInfos(ctx context.Context) []ml.DeviceInfo {
+	return nil
+}
+
+// HasExited returns true if the subprocess has exited.
 func (s *Server) HasExited() bool {
 	select {
 	case <-s.done: