GoMLX

https://gomlx.github.io/
Admin
GoMLX is an accelerated machine learning and math framework for Go that enables building and...
Tokens:59,610
Snippets:566
Trust Score:9.7
Update:1 month ago
Show doc for...
Tokens:
Raw
### Verify Go Installation

Source: https://gomlx.github.io/docs/first-model

Check the installed Go version to ensure it meets the minimum requirement of Go 1.26.

```bash
go version
```

--------------------------------

### Install and Use gomlx_checkpoints

Source: https://gomlx.github.io/docs/variables-and-checkpoints

Commands to install the GoMLX checkpoint utility and inspect a specific checkpoint directory.

```bash
# Build and install the tool
go install github.com/gomlx/gomlx/cmd/gomlx_checkpoints@latest

# Inspect a checkpoint folder
gomlx_checkpoints /path/to/checkpoints
```

--------------------------------

### Execution Logging Output Example

Source: https://gomlx.github.io/docs/debugging

Example output generated by the default and custom loggers.

```text
DefaultNodeLogger(Graph "Exec:main.SquareAndAddOne#1"):
	(Node #1) x^2 intermediate value: (Float32)[2]: []float32{4, 9}
Custom Logger: [Node #1] #full x^2 intermediate value = [25 49]
```

--------------------------------

### Sample IMDb dataset examples

Source: https://gomlx.github.io/gomlx/notebooks/imdb.html

Downloads the dataset if necessary and prints a specified number of random samples to the console.

```go
import "github.com/gomlx/gomlx/examples/imdb"

%%
AssertDownloaded()
imdb.PrintSample(3)
```

--------------------------------

### Settings File Format

Source: https://gomlx.github.io/docs/hyperparameters

Example content for a configuration file used with the -set file: prefix.

```text
# Settings file example
learning_rate=0.0005
batch_size=512
l2_regularization=1e-6
# Disable dropout for output
/output_layer/dropout=0.0
```

--------------------------------

### Training Output Logs

Source: https://gomlx.github.io/gomlx/notebooks/imdb.html

Example output logs showing data loading, backend initialization, and training progress metrics.

```text
> Loading previously generated preprocessed binary file.
Loaded data from "aclImdb.bin": 100000 examples, 141088 unique tokens, 23727054 tokens in total.
Backend "xla":	xla:cuda - PJRT "cuda" plugin (/home/janpf/.local/lib/go-xla/nvidia/pjrt_c_api_cuda_plugin.so) v0.100 [StableHLO] [1 device(s)]
Model: transformer
```

```text
        14% [====>...................................] (429 steps/s) [2s:9s] [step=724] [loss=1.16] [~loss=1.2] [~acc=49.66%]          
```

```text
       100% [========================================] (684 steps/s) [step=4999] [loss=0.625] [~loss=0.582] [~acc=91.01%]                
```

```text
	[Step 5000] median train step: 689 microseconds

Results on train-eval:
	Mean Loss (#loss): 0.608
	Mean Accuracy (#acc): 91.68%
Results on test-eval:
	Mean Loss (#loss): 0.676
	Mean Accuracy (#acc): 86.15%
```

--------------------------------

### Fetch GoMLX Dependencies

Source: https://gomlx.github.io/docs/first-model

Download the necessary GoMLX libraries and the MNIST example package.

```bash
go get github.com/gomlx/gomlx@latest
go get github.com/gomlx/compute@latest
go get github.com/gomlx/gomlx/examples/mnist@latest
```

--------------------------------

### Example output of dataset generation

Source: https://gomlx.github.io/gomlx/notebooks/dogsvscats.html

Sample console output showing the status of pre-generated binary files and image counts.

```text
Validation data for evaluation already generated in "/home/janpf/work/dogs_vs_cats/validation_eval_data.bin"
Training data for evaluation already generated in "/home/janpf/work/dogs_vs_cats/train_eval_data.bin"
Training data for training already generated in "/home/janpf/work/dogs_vs_cats/train_data.bin"

Dataset "train_data.bin": 1009400 images (== 50 x 20188)
	Elapsed time: 1m24.662268603s (83.873µs /image)
Dataset "train_eval_data.bin": 20188 images
Dataset "validation_eval_data.bin": 4798 images
```

--------------------------------

### Model Training Output Logs

Source: https://gomlx.github.io/gomlx/notebooks/imdb.html

Example logs showing data loading, backend initialization, training progress, and final evaluation metrics.

```text
> Loading previously generated preprocessed binary file.
Loaded data from "aclImdb.bin": 100000 examples, 141088 unique tokens, 23727054 tokens in total.
Backend "xla":	xla:cuda - PJRT "cuda" plugin (/home/janpf/.local/lib/go-xla/nvidia/pjrt_c_api_cuda_plugin.so) v0.100 [StableHLO] [1 device(s)]
Model: cnn
```

```text
         7% [=>......................................] (368 steps/s) [1s:25s] [step=719] [loss=1.58] [~loss=1.47] [~acc=51.61%]        
```

```text
       100% [========================================] (944 steps/s) [step=9999] [loss=0.368] [~loss=0.331] [~acc=96.32%]        %]        
```

```text
	[Step 10000] median train step: 560 microseconds

Results on train-eval:
	Mean Loss (#loss): 0.263
	Mean Accuracy (#acc): 98.90%
Results on test-eval:
	Mean Loss (#loss): 0.978
	Mean Accuracy (#acc): 84.07%
```

--------------------------------

### Import Pure Go Backend

Source: https://gomlx.github.io/docs/installation

Use the pure Go backend to avoid external dependencies or XLA installation.

```go
import _ "github.com/gomlx/gomlx/backends/simplego"
```

--------------------------------

### Training Loop Output

Source: https://gomlx.github.io/docs/get-started

Example console output generated during the training process.

```text
Starting training loop...
Step   999: MSE Loss = 0.000038 (moving average = 0.000040)
Step  1999: MSE Loss = 0.000067 (moving average = 0.000027)
Step  2999: MSE Loss = 0.000012 (moving average = 0.000030)
Step  3999: MSE Loss = 0.000010 (moving average = 0.000019)
Step  4999: MSE Loss = 0.000013 (moving average = 0.000018)
Training finished!
Successfully reconstructed image saved to reconstructed.png
```

--------------------------------

### Train a Multi-Layer Perceptron (MLP) in GoMLX

Source: https://gomlx.github.io/docs/get-started

A complete example showing how to prepare data, define a model architecture, and execute a training loop using the GoMLX trainer.

```go
// 1. Prepare training data: mapping (x, y) coordinates to (r, g, b) colors
inputs := make([][]float32, 0, width*height)
labels := make([][]float32, 0, width*height)

for y := range height {
	for x := range width {
		// Normalize coordinates to [-1.0, 1.0] for stable neural network training
		nx := float32(x)/float32(width)*2.0 - 1.0
		ny := float32(y)/float32(height)*2.0 - 1.0
		inputs = append(inputs, []float32{nx, ny})

		r, g, b, _ := img.At(bounds.Min.X+x, bounds.Min.Y+y).RGBA()
		labels = append(labels, []float32{float32(r) / 65535.0, float32(g) / 65535.0, float32(b) / 65535.0})
	}
}

backend := compute.MustNew()
store := model.NewStore()

// 2. Create an InMemoryDataset from the prepared coordinates and pixel colors
ds, err := dataset.InMemoryFromData(backend, "image_pixels", []any{inputs}, []any{labels})
if err != nil {
	log.Fatalf("failed to create dataset: %v", err)
}
// Configure dataset to yield random batches of size 512 continuously (infinitely)
ds.BatchSize(512, false).Shuffle().Infinite(true)

// 3. Define the neural network model function (MLP)
modelFn := func(scope *model.Scope, spec any, inputs []*Node) []*Node {
	x := inputs[0] // shape: [batch_size, 2]

	h := denseLayer(scope.In("layer1"), x, 64)
	h = activation.Relu(h)
	h = denseLayer(scope.In("layer2"), h, 64)
	h = activation.Relu(h)
	h = denseLayer(scope.In("layer3"), h, 64)
	h = activation.Relu(h)

	// Output RGB values mapped to [0.0, 1.0] using Sigmoid
	y := Sigmoid(denseLayer(scope.In("output"), h, 3))
	return []*Node{y}
}

// 4. Initialize Trainer with Adam optimizer and Mean Squared Error (MSE) loss
trainer := train.NewTrainer(
	backend,
	store,
	modelFn,
	loss.MeanSquaredError,
	optimizer.Adam().LearningRate(0.003).Done(),
	nil, // Train step metrics (optional)
	nil, // Eval metrics (optional)
)

// 5. Run the training loop for -steps (default 5000) steps
loop := train.NewLoop(trainer)
// Register a simple callback to print loss metrics periodically
train.EveryNSteps(loop, 1000, "log_metrics", 0, func(l *train.Loop, metrics []*tensors.Tensor) error {
	fmt.Printf("Step %5d: MSE Loss = %.6f (moving average = %.6f)\n", l.LoopStep, metrics[0].Value(), metrics[1].Value())
	return nil
})

fmt.Println("Starting training loop...")
_, err = loop.RunSteps(ds, *trainSteps)
if err != nil {
	log.Fatalf("training failed: %v", err)
}
fmt.Println("Training finished!")
```

--------------------------------

### Initialize a compute.Backend

Source: https://gomlx.github.io/docs/get-started

Create a backend at program startup to handle hardware execution. The system automatically selects the best available backend unless specified otherwise.

```go
import (
	"github.com/gomlx/compute"
	_ "github.com/gomlx/gomlx/backends/default" // Includes default backends.
)

backend, err := compute.New() // auto-selects best available backend
fmt.Printf("Backend: %s\n", backend.Description())
```

```text
Backend: xla:cuda - PJRT "cuda" plugin (/home/janpf/.local/lib/go-xla/nvidia/pjrt_c_api_cuda_plugin.so) v0.112 [StableHLO] [1 device(s)]
```

--------------------------------

### Initialize Go Project

Source: https://gomlx.github.io/docs/first-model

Create a new directory and initialize a Go module for the project.

```bash
mkdir gomlx-mnist-quickstart
cd gomlx-mnist-quickstart
go mod init gomlx-mnist-quickstart
```

--------------------------------

### NanLogger Output Example

Source: https://gomlx.github.io/docs/debugging

Example of the log output generated when a NaN or Inf value is detected.

```text
I0720 18:14:08.141136  135770 nanlogger.go:375] NaN/Inf observed for scope ["dense_layer_1" "output_division"]
```

--------------------------------

### Disable Automatic XLA Installation

Source: https://gomlx.github.io/docs/installation

Programmatically prevent the automatic installation of XLA PJRT plugins.

```go
xla.EnableAutoInstall(false)
```

--------------------------------

### Configure and execute a training loop

Source: https://gomlx.github.io/docs/first-model

Initializes the backend, prepares datasets, and runs a training loop with progress tracking and evaluation reporting.

```go
backend := compute.MustNew()
defer backend.Finalize()
fmt.Printf("Backend: %s (%s)\n", backend.Name(), backend.Description())

store := model.NewStore()
trainDS, testDS := prepareDatasets(backend, dataDir)

accuracyMetric := metric.NewSparseCategoricalAccuracy("Accuracy", "acc")
trainer := train.NewTrainer(
	backend,
	store,
	ConvModel,
	loss.SparseCategoricalCrossEntropyLogits, // cross-entropy loss for integer labels
	optimizer.Adam().LearningRate(1e-3).Done(),
	[]metric.Interface{accuracyMetric}, // metrics reported during training
	[]metric.Interface{accuracyMetric}, // metrics reported during evaluation
)

const epochs = 5
loop := train.NewLoop(trainer)
// AttachProgressBar gives you a live progress bar with loss/metric values -- there's no
// need to print them by hand.
commandline.AttachProgressBar(loop)

fmt.Printf("Starting training for %d epochs...\n", epochs)
if _, err := loop.RunEpochs(trainDS, epochs); err != nil {
	log.Fatalf("Training loop failed: %+v", err)
}

fmt.Println("\nEvaluating model performance on test set...")
if err := commandline.ReportEval(trainer, testDS); err != nil {
	log.Fatalf("Evaluation failed: %+v", err)
}
```

--------------------------------

### Initialize and print model settings

Source: https://gomlx.github.io/gomlx/notebooks/dogsvscats.html

Initializes the store from settings and prints the current configuration using commandline utilities.

```go
%%
store, _ := StoreFromSettings()
fmt.Println(commandline.SprintSettings(store))
```

--------------------------------

### Configure and Initialize a Trainer

Source: https://gomlx.github.io/docs/training-loop

Sets up the backend, model store, model function, metrics, and trainer instance for training.

```go
backend := compute.MustNew()
store := model.NewStore()

// Configure hyperparameters in the model store
store.SetParams(map[string]any{
	"batch_size":                128,
	"train_steps":               1000,
	optimizer.ParamOptimizer:    "adam",
	optimizer.ParamLearningRate: 0.001,
	activation.ParamActivation:  "relu",
	fnn.ParamNumHiddenLayers:    2,
	fnn.ParamNumHiddenNodes:     64,
})
scope := store.RootScope()
modelFn := func(scope *model.Scope, spec any, inputs []*Node) []*Node {
	// In a real example, inputs would contain calibrated continuous and categorical features.
	// For simplicity, we just run a Feed-Forward Neural Network on the continuous features.
	continuousFeatures := inputs[1]
	logits := fnn.New(scope, continuousFeatures, 1).Done()
	return []*Node{logits}
}

// 4. Configure metrics
// Moving average is recommended for training since the model changes at each step.
movingAccuracy := metric.NewMovingAverageBinaryLogitsAccuracy("Moving Average Accuracy", "~acc", 0.01)
// Mean accuracy is recommended for evaluation because the model is frozen.
meanAccuracy := metric.NewMeanBinaryLogitsAccuracy("Mean Accuracy", "#acc")

// 5. Create the Trainer
trainer := train.NewTrainer(
	backend,
	store,
	modelFn,
	loss.BinaryCrossentropyLogits,
	optimizer.Adam().LearningRate(0.001).Done(),
	[]metric.Interface{movingAccuracy}, // metrics evaluated at train time
	[]metric.Interface{meanAccuracy},   // metrics evaluated at eval time
)
```

--------------------------------

### Run MNIST tutorial via go run

Source: https://gomlx.github.io/docs/first-model

Executes the tutorial directly from the repository using the latest version.

```bash
go run github.com/gomlx/gomlx/examples/gomlx.github.io/mnist/first-model@latest
```

--------------------------------

### Dataset memory output

Source: https://gomlx.github.io/gomlx/notebooks/flow_matching.html

Example output showing the memory footprint of the loaded datasets.

```text
Total number of examples: #train=6487, #validation=1702
trainDS (in-memory) using 304.2 MiB of memory.
validationDS (in-memory) using 79.8 MiB of memory.
```

--------------------------------

### Run Training from Command Line

Source: https://gomlx.github.io/gomlx/notebooks/flow_matching.html

Commands to initiate training for 64x64 and 128x128 image resolutions from the demo directory.

```bash
go run . --checkpoint fm_64x64
```

```bash
go run . --checkpoint=fm_128x128 --set='image_size=128;diffusion_channels_list=32,64,96,128,160;diffusion_num_residual_blocks=6;learning_rate=1e-2;diffusion_loss=huber'
```

--------------------------------

### Configure and Download Dogs vs Cats Dataset

Source: https://gomlx.github.io/gomlx/notebooks/dogsvscats.html

Initializes the dataset store, parses command-line flags for configuration, and triggers the download of the dataset.

```go
import (
    "github.com/gomlx/gomlx/ml/dataset"
    "github.com/gomlx/gomlx/ml/model"
    "github.com/gomlx/gomlx/examples/dogsvscats"
    "github.com/janpfeifer/must"
)

var (
    flagDataDir   = flag.String("data", "~/work/dogs_vs_cats", "Directory to cache downloaded and generated dataset files.")
    flagCheckpoint = flag.String("checkpoint", "", "Directory save and load checkpoints from. If left empty, no checkpoints are created.")
    flagSet *string
)

// init_scopeSettings is executed at initialization and sets the flag "set" to accept the default scope setting flags.
// Use --help to see all options one can set.
func init_scopeSettings() {
	store := dogsvscats.CreateModelStore()
	flagSet = commandline.CreateSettingsFlag(store, "set")
}

// StoreFromSettings returns the default scope merged with values parsed from -set flag.
// It also returns the list of parameters set with -set flag.
func StoreFromSettings() (store *model.Store, paramsSet []string) {
    store = dogsvscats.CreateModelStore()
    paramsSet = must.M1(commandline.ParseSettings(store, *flagSet))
    return
}

%%
// Create Store and parse -set flag.
store, _ := StoreFromSettings()
config := dogsvscats.NewPreprocessingConfiguration(store, *flagDataDir)

// Download dataset, if not yet downloaded.
must.M(dogsvscats.Download(config.DataDir))
fmt.Println("Dogs vs Cats dataset downloaded.")
```

```text
Dogs vs Cats dataset downloaded.
```

--------------------------------

### Download and List Files from HuggingFace Hub

Source: https://gomlx.github.io/docs/go-huggingface

Demonstrates initializing a repository connection, iterating through file information, and downloading a specific file using the hub package.

```go
import (
	"fmt"
	"os"

	"github.com/gomlx/compute/support/humanize"
	"github.com/gomlx/go-huggingface/hub"
)

func main() {
	hfAuthToken := os.Getenv("HF_TOKEN") // Optional token for private or gated models
	repo := hub.New("google/gemma-2-2b-it").WithAuth(hfAuthToken)

	// Iterate over files in the repository
	for fileInfo, err := range repo.IterFileInfos() {
		if err != nil {
			panic(err)
		}
		fmt.Printf("File: %s (Size: %s)\n", fileInfo.Name, humanize.Bytes(fileInfo.Size))
	}

	// Download a specific file to the cache
	localPath, err := repo.DownloadFile("config.json")
	if err != nil {
		panic(err)
	}
	fmt.Printf("Downloaded config.json to: %s\n", localPath)
}
```

--------------------------------

### Euclidean distance output

Source: https://gomlx.github.io/docs/get-started

The expected console output from the Euclidean distance computation example.

```text
Distance (1,2)->(4,6):  float32(5)
Distance (0,0)->(5,12): float32(13)
```

--------------------------------

### Initialize Imports and Configuration

Source: https://gomlx.github.io/gomlx/notebooks/flow_matching.html

Sets up the necessary GoMLX packages, flags for data management, and the configuration function for the Flow Matching model.

```go
import (
    "github.com/gomlx/compute"
    fm "github.com/gomlx/gomlx/examples/FlowMatching"
    flowers "github.com/gomlx/gomlx/examples/oxfordflowers102"    
    "github.com/gomlx/gomlx/examples/oxfordflowers102/diffusion"    
    . "github.com/gomlx/gomlx/core/graph"
    "github.com/gomlx/gomlx/ml/model"
    "github.com/gomlx/gomlx/ml/train"
    timage "github.com/gomlx/gomlx/core/tensors/images"
    "github.com/janpfeifer/gonb/gonbui"
    "github.com/janpfeifer/must"

    _ "github.com/gomlx/gomlx/backends/default"
)

type Dataset = train.Dataset

var (
    backend = compute.MustNew()
    _ *Node = nil

    flagDataDir    = flag.String("data", "~/work/oxfordflowers102", "Directory to cache downloaded and generated dataset files.")
	flagEval       = flag.Bool("eval", true, "Whether to evaluate the model on the validation data in the end.")
	flagVerbosity  = flag.Int("verbosity", 1, "Level of verbosity, the higher the more verbose.")
	flagCheckpoint = flag.String("checkpoint", "", "Directory save and load checkpoints from. If left empty, no checkpoints are created.")

    // settings is bound to a "-set" flag to be used to set scope hyperparameters.
    settings = commandline.CreateSettingsFlag(fm.CreateStore(), "set")
)

func NewConfig() (config *diffusion.Config) {
    store := fm.CreateStore()
    paramsSet := must.M1(commandline.ParseSettings(store, *settings))
    config = diffusion.NewConfig(backend, store, *flagDataDir, paramsSet)
    if *flagCheckpoint != "" {
        config.AttachCheckpoint(*flagCheckpoint)
    }
    return
}

%%
cfg := NewConfig()
must.M(flowers.DownloadAndParse(cfg.DataDir))
fmt.Println("Oxford Flowers 102 dataset downloaded:")
fmt.Printf("\t%d images, %d labels, %d examples\n", len(flowers.AllImages), len(flowers.AllLabels), flowers.NumExamples)
```

--------------------------------

### Clean Training Directory

Source: https://gomlx.github.io/gomlx/notebooks/flow_matching.html

Removes existing checkpoints before starting a new training run.

```bash
!rm -rf ~/work/oxfordflowers102/fm_base
```

--------------------------------

### Initialize Imports and Global Variables

Source: https://gomlx.github.io/gomlx/notebooks/flow_matching.html

Sets up the necessary imports for GoMLX, Gonum plotting, and initializes the default compute backend.

```go
import (
    "bytes"
    "flag"
    colors "image/color"
    "github.com/gomlx/compute"
    _ "github.com/gomlx/gomlx/backends/default"
    fm "github.com/gomlx/gomlx/examples/FlowMatching"
    . "github.com/gomlx/gomlx/core/graph"
    "github.com/gomlx/gomlx/ml/model"
    "github.com/janpfeifer/gonb/gonbui"
    "github.com/janpfeifer/must"
    "gonum.org/v1/plot"
    "gonum.org/v1/plot/plotter"
    plotvg "gonum.org/v1/plot/vg"
)

var (
    backend = compute.MustNew()
    _ *Node = nil
)
```

--------------------------------

### Override Parameters via CLI

Source: https://gomlx.github.io/docs/hyperparameters

Examples of using the -set flag for direct, scoped, and file-based hyperparameter overrides.

```bash
go run ./cmd/train -set "learning_rate=0.005;batch_size=256"
```

```bash
go run ./cmd/train -set "/output_layer/l2_regularization=0.0;learning_rate=0.01"
```

```bash
go run ./cmd/train -set "file:configs/hparams_sweep_3.txt"
```

--------------------------------

### Execute Short Training Run

Source: https://gomlx.github.io/gomlx/notebooks/dogsvscats.html

Runs a 50-step training session with default settings to verify the environment and model configuration.

```go
%% --set="model=cnn;plots=false;train_steps=50"
store, paramsSet := StoreFromSettings()
dogsvscats.Train(store, *flagDataDir, *flagCheckpoint, /* eval */ false, paramsSet)
```

--------------------------------

### Manual Dataset Consumption

Source: https://gomlx.github.io/docs/datasets

Example of manually iterating over a dataset and explicitly freeing device memory using Finalize.

```go
// Manual dataset consumption example
next, stop := iter.Pull2(myDataset.Iter())
defer stop()

for {
	batch, err, ok := next()
	if !ok {
		break
	}
	if err != nil {
		log.Fatalf("failed reading: %v", err)
	}

	// 1. Process batch...
	doSomething(batch.Inputs)

	// 2. Mandatory: free device memory
	batch.Finalize()
}
```

--------------------------------

### Import Checkpoint Packages

Source: https://gomlx.github.io/docs/first-model

Include the necessary time and checkpoint packages to enable model persistence.

```go
"time"
"github.com/gomlx/gomlx/ml/model/checkpoint"
```

--------------------------------

### Calculate Euclidean distance with GoMLX

Source: https://gomlx.github.io/docs/get-started

A minimal example demonstrating a stateless computation using a graph to calculate Euclidean distance between two tensors.

```go
func MyComputation(x, y *Node) *Node {
	return Sqrt(ReduceAllSum(Square(Sub(x, y))))
}

backend, err := compute.New()
exec, err := NewExec(backend, MyComputation)
results, err := exec.Call([]float32{1.0, 2.0}, []float32{4.0, 6.0})
fmt.Printf("Distance (1,2)->(4,6):  %v\n", results[0])
results, err = exec.Call([]float32{0, 0}, []float32{5, 12})
fmt.Printf("Distance (0,0)->(5,12): %v\n", results[0])
```

--------------------------------

### Example Output of Dataset Batch

Source: https://gomlx.github.io/gomlx/notebooks/uci-adult.html

This block shows the console output generated by the dataset processing logic, including tensor shapes and label distributions.

```text
Backend: xla - xla:cuda - PJRT "cuda" plugin (/home/janpf/.local/lib/go-xla/nvidia/pjrt_c_api_cuda_plugin.so) v0.100 [StableHLO] [1 device(s)]
Inputs of batch (size 128):
	categorical:
		Features=[workclass education marital-status occupation relationship race sex native-country]
		Values: [128][8]int64{
 {4, 16, 3, ..., 5, 2, 39},
 {6, 8, 3, ..., 5, 2, 0},
 {4, 10, 5, ..., 5, 1, 39},
 ...,
 {4, 6, 3, ..., 5, 2, 39},
 {4, 8, 3, ..., 2, 2, 39},
 {4, 12, 1, ..., 5, 1, 39}}
	continuous:
		Features=[age education-num capital-gain capital-loss hours-per-week]
		Values: [128][5]float32{
 {58, 10, 0, 0, 36},
 {42, 12, 0, 0, 40},
 {46, 13, 0, 0, 40},
 ...,
 {55, 4, 0, 0, 75},
 {41, 12, 0, 0, 40},
 {36, 9, 0, 0, 40}}
	weights: [128][1]float32{
 {9.584e+04},
 {1.838e+05},
 {2.792e+05},
 ...,
 {1.905e+05},
 {3.497e+05},
 {1.166e+05}}

Labels of batch:
	[128][1]float32{
 {0},
 {0},
 {0},
 ...,
 {0},
 {0},
 {0}}

Labels distributions:
	Train:	24.08% positive
	Test:	23.62% positive
```

--------------------------------

### Implement Custom Mean Squared Error

Source: https://gomlx.github.io/docs/losses

Example of a custom loss function calculating the mean squared error between true labels and model predictions.

```go
import . "github.com/gomlx/gomlx/core/graph"

func CustomMSE(labels, predictions []*Node) *Node {
	yTrue := labels[0]
	yPred := predictions[0]
	
	// Compute mean squared error: Mean((yTrue - yPred)^2)
	squaredDiff := Square(Sub(yTrue, yPred))
	return ReduceMean(squaredDiff)
}
```

--------------------------------

### Apply Activation Functions

Source: https://gomlx.github.io/docs/layers

Demonstrates applying a specific ReLU activation and using the scope-based configuration helper.

```go
import "github.com/gomlx/gomlx/ml/layers/activation"

// RELU activation.
h = activation.Relu(h)

// Apply activation configured (hyperparameter) in current scope.
h = activation.ApplyFromScope(scope, h)
```

--------------------------------

### Execute Training Loop

Source: https://gomlx.github.io/docs/training-loop

Initiates the training process by running the loop until a specified global step is reached, handling potential restarts.

```go
trainSteps := model.GetParamOr(scope, "train_steps", 1000)
globalStep := int(optimizer.GetGlobalStep(scope))
if globalStep != 0 {
	fmt.Printf("Restarting training from global step %d\n", globalStep)
}

_, err := loop.RunToGlobalStep(trainDS, trainSteps)
if err != nil {
	panic(err)
}
fmt.Println("Training complete!")
```

--------------------------------

### Sample and display images from the dataset

Source: https://gomlx.github.io/gomlx/notebooks/dogsvscats.html

Defines the sampling function and helper to render images as HTML in a notebook. Requires a configured preprocessing configuration.

```go
import(
    timage "github.com/gomlx/gomlx/core/tensors/images"
    "github.com/gomlx/gomlx/core/tensors"
    "github.com/gomlx/gomlx/ml/dataset"
    "github.com/gomlx/compute/dtypes"

    "github.com/janpfeifer/gonb/gonbui"
)

// sample some random augmented images and display them in the Notebook.
func sample(config *dogsvscats.PreprocessingConfiguration, numRows, numPerRow int) {
    var images []image.Image
    var labels []dogsvscats.DogOrCat

    // Create datasets with batch size equals 2*numPerRow.
    var configForSample dogsvscats.PreprocessingConfiguration
    configForSample = *config
    configForSample.BatchSize = numRows*numPerRow  // Sample only what we need, in one batch.
    configForSample.ModelImageSize = 256
    configForSample.ForceOriginal = true
    configForSample.UseParallelism = true
    configForSample.DType = dtypes.Uint8
    
    // Sample the images.
    ds, _, _ := dogsvscats.CreateDatasets(&configForSample)
    var batch train.Batch
    for b, err := range ds.Iter() {
        must.M(err)
        batch = b
        break
    }
    inputsT, labelsT := batch.Inputs, batch.Labels
    
    // Get indices and labels of the images.
    indices := inputsT[1].Value().([]int64)
    labelsFloat := labelsT[0].Value().([]uint8)
    labels = make([]dogsvscats.DogOrCat, 0, numRows*numPerRow)
    for _, labelFloat := range labelsFloat {
        labels = append(labels, dogsvscats.DogOrCat(labelFloat))
    }
    
    // Convert images from tensor to Go images.
    images = timage.ToImage().Batch(inputsT[0])
    htmlRows := make([]string, 0, numRows)
    count := 0
    for row := 0; row < numRows; row++ {
        cells := make([]string, 0, numPerRow)
        for col := 0; col < numPerRow; col++ {
            imgIdx := indices[count]
            cells = append(cells, embedImageInHTML(images[count], labels[count].String(), imgIdx, 256))
            count++
        }
        htmlRows = append(htmlRows, fmt.Sprintf("<tr>\n\t<td>%s</td>\n</tr>", strings.Join(cells, "</td>\n\t<td>")))
    }
    htmlTable := fmt.Sprintf("<h3>%s</h3><table>%s</table>\n", "Sample Dogs vs Cats", strings.Join(htmlRows, ""))
    gonbui.DisplayHTML(htmlTable)
}

// embedImageInHTML, with a label.
func embedImageInHTML(img image.Image, label string, imgIdx int64, size int) string {
    imgSrc := must.M1(gonbui.EmbedImageAsPNGSrc(img))   // Generate image in-html (in the src filed), as opposed to a separate file.
    return fmt.Sprintf(`<figure style="padding:4px;text-align: center; background-color: lightgray; color: black;"><img src="%s" width="%dpx" height="%dpx"><figcaption style="text-align: center;">%s (%d)</figcaption></figure>`,
                       imgSrc, size, size, label, imgIdx)
}

%%
store, _ := StoreFromSettings()
config := dogsvscats.NewPreprocessingConfiguration(store, *flagDataDir)
sample(config, 2, 8)
```

--------------------------------

### Configuring Dataset Iteration Modes

Source: https://gomlx.github.io/docs/datasets

Demonstrates configuring batch size, shuffling, and infinite looping for training and evaluation. Uses .Copy() to reuse underlying data across different configurations.

```go
// A. Create base in-memory dataset for testing --the default configuration
//    has no shuffling and yields one epoch only.
testEvalDS  := adult.NewDataset(backend, adult.Data.Test, "test").BatchSize(batchSize, false)

// B. Eval on training: no shuffling and one epoch.
baseTrainDS := adult.NewDataset(backend, adult.Data.Train, "batched train")
trainEvalDS := baseTrainDS.BatchSize(batchSize, false)

// C. Copy the underlying data (shallow, only references are copied), and configure
//    to loop indefinitely and shuffling for training. 
trainDS := baseTrainDS.BatchSize(batchSize, true).Shuffle().Infinite(true)
```

--------------------------------

### Initialize a Training Loop

Source: https://gomlx.github.io/docs/training-loop

Create a new training loop instance using a trainer object.

```go
loop := train.NewLoop(trainer)
```

--------------------------------

### Override Hyperparameters via Command Line

Source: https://gomlx.github.io/gomlx/notebooks/uci-adult.html

Demonstrates how to override default hyperparameters using the -set flag and verify the changes.

```go
// Let's test that we can set hyperparameters by setting it in the "-set" flag:
%% -set="batch_size=17;fnn_num_hidden_layers=12"
store, paramsSet := storeFromSettings()
for _, name := range paramsSet {
    if value, found := store.GetParam(name); found {
        fmt.Printf("\t%s=%v\n", name, value)
    }
}
```

--------------------------------

### Create in-memory datasets

Source: https://gomlx.github.io/gomlx/notebooks/flow_matching.html

Initializes in-memory datasets and displays memory usage statistics for the Oxford 102 Flowers dataset.

```go
// Remove cached file to force regeneratrion.
// !rm -f "${HOME}/work/oxfordflowers102/"*_cached_images_*
import "github.com/gomlx/gomlx/support/fsutil"

%% --set="image_size=128"
cfg := NewConfig()
trainDS, validationDS := cfg.CreateInMemoryDatasets()

fmt.Println()
fmt.Printf("Total number of examples: #train=%d, #validation=%d\n", trainDS.NumExamples(), validationDS.NumExamples())
fmt.Printf("trainDS (in-memory) using %s of memory.\n", humanize.Bytes(trainDS.ByteSize()))
fmt.Printf("validationDS (in-memory) using %s of memory.\n", humanize.Bytes(validationDS.ByteSize()))

// Output a random sample.
trainDS.Shuffle()
sampleTable("Oxford 102 Flowers Sample -- In-Memory Dataset", trainDS, 1, 4)
```

--------------------------------

### Configure Feed-Forward Networks

Source: https://gomlx.github.io/docs/layers

Demonstrates creating standard MLPs, ensembles, and scope-configured FNNs using the fnn package.

```go
import "github.com/gomlx/gomlx/ml/layers/fnn"

// Standard MLP configured fluently
y := fnn.New(scope, input, outputDim).
    NumHiddenLayers(3, 256).
    Activation(activation.Swish).
    Normalization("layer").
    Dropout(0.1).
    Done()

// Ensemble of 5 MLPs running in parallel
ensembleY := fnn.New(scope, input, outputDim).
    NumHiddenLayers(2, 64).
    WithEnsembleSize(5).
    Done()

// Automatically configured FNN, all (num hidden layers, num hidden nodes, activation, etc.) configured by scope hyperparameters.
logits = fnn.New(scope.In("fnn"), logits, 1).Done()
```

--------------------------------

### Train Model in Notebook

Source: https://gomlx.github.io/gomlx/notebooks/flow_matching.html

Executes a short training run for development testing purposes.

```go
%% --checkpoint fm_base --set=train_steps=100;plots=false
cfg := NewConfig()
fm.TrainModel(cfg, *flagCheckpoint, *flagEval, *flagVerbosity)
```

--------------------------------

### Implement Execution Logging

Source: https://gomlx.github.io/docs/debugging

Demonstrates marking nodes for logging during graph construction and configuring both default and custom loggers.

```go
func SquareAndAddOne(x *Node) *Node {
	xSquared := Mul(x, x)
	// Mark the intermediate square calculation to be logged:
	xSquared.SetLoggedf("#full x^2 intermediate value")

	one := Scalar(x.Graph(), x.DType(), 1.0)
	return Add(xSquared, one)
}

exec := MustNewExec(backend, SquareAndAddOne)

// Option 1: Run with the default logger (prints to stdout)
exec.MustCall([]float32{2.0, 3.0})

// Option 2: Run with a custom logger
exec.SetNodeLogger(func(g *Graph, messages []string, values []*tensors.Tensor, nodes []NodeId) {
	for i, msg := range messages {
		fmt.Printf("Custom Logger: [Node #%d] %s = %v\n", nodes[i], msg, values[i].Value())
	}
})
exec.MustCall([]float32{5.0, 7.0})
```

--------------------------------

### Configure Model Store with Command Line Flags

Source: https://gomlx.github.io/gomlx/notebooks/imdb.html

Uses a command line flag to initialize and update model hyperparameters via the model.Store object.

```go
import (
    "golang.org/x/exp/maps"
    "github.com/gomlx/gomlx/ml/model"
)

// settings is bound to a "-set" flag to be used to set store hyperparameters.
var settings = commandline.CreateSettingsFlag(imdb.CreateModelStore(), "set")

// StoreFromSettings is the default store (createModelStore) changed by -set flag.
// It also returns the list of parameters changed by -set in paramsSet: we use this later to avoid loading over the values from checkpoint.
func StoreFromSettings() (store *model.Store, paramsSet []string) {
    store = imdb.CreateModelStore()
    paramsSet = must.M1(commandline.ParseSettings(store, *settings))
    return store, paramsSet
}

%% -set="model=cnn"
fmt.Printf("Model types: %q\n", maps.Keys(imdb.ValidModels))
store, _ := StoreFromSettings()
fmt.Println(commandline.SprintSettings(store.RootScope()))
```

--------------------------------

### Train InceptionV3 with fine-tuning

Source: https://gomlx.github.io/gomlx/notebooks/dogsvscats.html

Executes the training process for InceptionV3 with pre-trained weights and fine-tuning enabled for 2000 steps.

```go
%% --checkpoint=inceptionv3_finetuned --set="model=inception;inception_pretrained=true;inception_finetuning=true;plots=true;train_steps=2000"
store, paramsSet := StoreFromSettings()
dogsvscats.Train(store, *flagDataDir, *flagCheckpoint, true, paramsSet)
```

--------------------------------

### Define Model Hyperparameters with model.Store

Source: https://gomlx.github.io/gomlx/notebooks/uci-adult.html

Initializes a model.Store with default hyperparameters for training, optimization, and network architecture.

```go
import (
    "github.com/gomlx/gomlx/ml/model"
    "github.com/gomlx/gomlx/ml/layers"
    "github.com/gomlx/gomlx/ml/layers/fnn"
    "github.com/gomlx/gomlx/ml/layers/kan"
    "github.com/gomlx/gomlx/ml/layers/regularizer"
    "github.com/gomlx/gomlx/ui/commandline"
    "github.com/gomlx/gomlx/ml/train/optimizer"
    "github.com/gomlx/gomlx/ml/train/optimizer/cosineschedule"
	"github.com/janpfeifer/must"    
)

// settings is bound to a "-set" flag to be used to set scope hyperparameters.
var settings = commandline.CreateSettingsFlag(createModelStore(), "set")

func createModelStore() *model.Store {
	store := model.NewStore()
	store.SetParams(map[string]any{
		"train_steps":     5000,
		"batch_size":      128,
		"eval_batch_size":      1000,
		"plots":           true,
        "num_checkpoints": 3,

		optimizer.ParamOptimizer:           "adam",
		optimizer.ParamLearningRate:        0.001,
		optimizer.ParamAdamEpsilon:         1e-7,
		optimizer.ParamAdamDType:           "",
		cosineschedule.ParamPeriodSteps:     0,
		activation.ParamActivation:         "sigmoid",
		layers.ParamDropoutRate:             0.0,
		regularizer.ParamL2:                1e-5,
		regularizer.ParamL1:                1e-5,

		// FNN network parameters:
		fnn.ParamNumHiddenLayers: 1,
		fnn.ParamNumHiddenNodes:  4,
		fnn.ParamResidual:        true,
		fnn.ParamNormalization:   "layer",

		// KAN network parameters:
		"kan":                       false, // Enable kan
		kan.ParamNumControlPoints:   20,    // Number of control points
		kan.ParamNumHiddenNodes:     4,
		kan.ParamNumHiddenLayers:    1,
		kan.ParamBSplineDegree:      2,
		kan.ParamBSplineMagnitudeL1: 1e-5,
		kan.ParamBSplineMagnitudeL2: 0.0,
		kan.ParamDiscrete:           false,
		kan.ParamDiscreteSoftness:   0.1,
	})
	return store
}

// storeFromSettings is the default scope (createModelStore) changed by -set flag.
func storeFromSettings() (store *model.Store, paramsSet []string) {
    store = createModelStore()
    paramsSet = must.M1(commandline.ParseSettings(store, *settings))
    return
}
```

--------------------------------

### Define Datasets and Positive Ratio Calculation

Source: https://gomlx.github.io/gomlx/notebooks/uci-adult.html

This snippet initializes the backend, defines the BuildDatasets function to create training and evaluation samplers, and implements a helper to calculate the positive label ratio.

```go
import (
    "flag"
    "fmt"
    "io"

    "github.com/gomlx/compute"
    "github.com/gomlx/gomlx/examples/adult"
    "github.com/gomlx/gomlx/ml/train"
    "github.com/gomlx/gomlx/core/tensors"
    . "github.com/gomlx/gomlx/core/graph"

    _ "github.com/gomlx/gomlx/backends/default"
)

// Global backend created an initialization, used everywhere.
var backend = compute.MustNew()

// BuildDatasets returns 3 `train.Dataset`:
// * trainDS is an endless random sampler used for training.
// * trainEvalDS samples through exactly one epoch of the train dataset.
// * testEvalDS samples through exactly one epoch of the test dataset.
func BuildDatasets(store *model.Store) (trainDS, trainEvalDS, testEvalDS train.Dataset) {
	batchSize := model.GetRootParamOr(store, "batch_size", 128)
	evalBatchSize := model.GetRootParamOr(store, "eval_batch_size", 1000)
    baseDS := adult.NewDataset(backend, adult.Data.Train, "batched train")
    trainEvalDS = baseDS.Copy().BatchSize(evalBatchSize, false)
    testEvalDS = adult.NewDataset(backend, adult.Data.Test, "test").
        BatchSize(evalBatchSize, false)
    // For training, we shuffle and loop indefinitely.
    trainDS = baseDS.BatchSize(batchSize, true).Shuffle().Infinite(true)
    return
}

// PositiveRatio finds out the the ratio of positive labels in the
// training and testing data.
//
// We could do this easily with GoMLX computation model (just `ReduceAllSum`), but
// this examples shows it's also ok to mix Go computations.
func PositiveRatio(ds train.Dataset) float32 {
    var sum float32
    var count float32
    reduceSumExec := MustNewExec1(backend, func(labels *Node) *Node {
        return ReduceAllSum(labels)
    })
    for batch, err := range ds.Iter() {
        if err != nil { panic(err) }
        sum += reduceSumExec.MustCall(batch.Labels[0]).Value().(float32)
        count += float32(batch.Labels[0].Size())
        batch.Finalize()
    }
    return sum/count
}

%%
fmt.Printf("Backend: %s - %s\n", backend, backend.Description())

adult.LoadAndPreprocessData(*flagDataDir, *flagNumQuantiles, *flagForceDownload, *flagVerbosity)    
store, _ := storeFromSettings()
trainDS, trainEvalDS, testEvalDS := BuildDatasets(store)

// Take one batch.
var batch train.Batch
var err error
for batch, err = range trainDS.Iter() {
    break
}
if err != nil { panic(err) }
fmt.Printf("Inputs of batch (size %d):\n", model.GetRootParamOr(store, "batch_size", 0))
fmt.Printf("\tcategorical:\n\t\tFeatures=%v\n", adult.Data.VocabulariesFeatures)
fmt.Printf("\t\tValues: %s\n", batch.Inputs[0])
fmt.Printf("\tcontinuous:\n\t\tFeatures=%v\n", adult.Data.QuantilesFeatures)
fmt.Printf("\t\tValues: %s\n", batch.Inputs[1])
fmt.Printf("\tweights: %s\n", batch.Inputs[2])
fmt.Printf("\nLabels of batch:\n\t%s\n", batch.Labels[0])
fmt.Printf("\nLabels distributions:\n\tTrain:\t%.2f%% positive\n\tTest:\t%.2f%% positive\n",
           PositiveRatio(trainEvalDS)*100.0, PositiveRatio(testEvalDS)*100.0)
```

--------------------------------

### Sample Source and Target Distributions

Source: https://gomlx.github.io/gomlx/notebooks/flow_matching.html

Generates points from a 2D Gaussian source distribution and a 'two moons' target distribution for visualization.

```go
numPoints := 200
store := model.NewStore()
normalPoints := model.MustCallOnce(backend, store, func (scope *model.Scope, g *Graph) *Node {
        return scope.RandomNormal(g, shapes.Make(dtypes.F32, numPoints, 2))
    }).Value().([][]float32)
moonsPoints := model.MustCallOnce(backend, store, func (scope *model.Scope, g *Graph) *Node {
        return fm.MakeMoons(scope, g, numPoints)
    }).Value().([][]float32)
gonbui.DisplayHTMLF("<table><tr><td>%s</td><td>%s</td></tr></table>",
    HistogramXYs("Source Distribution: Normal", normalPoints, 200, 200),
    HistogramXYs("Target Distribution: Moons", moonsPoints, 200, 200))
```
Tokens:
Raw