!*rm -f go.work && go work init && go work use . "${HOME}/Projects/gomlx"
%goworkfix

	- Added replace rule for module "github.com/gomlx/gomlx" to local directory "/home/janpf/Projects/gomlx".

import . "github.com/gomlx/gomlx/graph"

func SumGraph(a, b *Node) *Node {
    return Add(a, b)
}

import _ "github.com/gomlx/gomlx/backends/xla"

var backend = backends.New()

%%

// Short version: `ExecOnce` will compile the function and execute only once.
two := ExecOnce(backend, /* GoMLX function */ SumGraph, /* Arguments */ 1.0, 1.0)
fmt.Printf("Short version:\t1+1=%v\n", two)

// Exec version: `NewExec` will compile the SumGraph function so it can be executed efficiently
// many times.
exec := NewExec(backend, SumGraph)  
two = exec.Call(1, 1)[0]
fmt.Printf("Exec version:\t1+1=%v\n", two.Value())

Short version:	1+1=float64(2.0000)
Exec version:	1+1=2

import (
    "fmt"
    . "github.com/gomlx/gomlx/graph"
)

func SumGraph(a, b *Node) *Node {
    fmt.Printf("\t. building graph for a.shape=%s and b.shape=%s\n", a.Shape(), b.Shape())
    return Add(a, b)
}


func main() {
    sumExec := NewExec(backend, SumGraph)
    two := sumExec.Call(1, 1)[0]
    fmt.Printf("1+1=%v\n", two.Value())

    for ii := 0; ii < 5; ii++ {
        sumInts := sumExec.Call(ii, ii)[0]
        fmt.Printf("%d+%d=%v\n", ii, ii, sumInts.Value())
    }

    five := sumExec.Call(3.5, 1.5)[0]
    fmt.Printf("3.5+1.5=%v\n", five.Value())

    many := sumExec.Call([]float32{1.1, 2.2, 3.3}, []float32{10, 10, 10})[0]
    fmt.Printf("[1.1, 2.2, 3.3] + [10, 10, 10] = %v\n", many.Value())
}

	. building graph for a.shape=(Int64) and b.shape=(Int64)
1+1=2
0+0=0
1+1=2
2+2=4
3+3=6
4+4=8
	. building graph for a.shape=(Float64) and b.shape=(Float64)
3.5+1.5=5
	. building graph for a.shape=(Float32)[3] and b.shape=(Float32)[3]
[1.1, 2.2, 3.3] + [10, 10, 10] = [11.1 12.2 13.3]

%%
sumExec := NewExec(backend, SumGraph)
_ = sumExec.Call(1.1, 2) // Error: arguments have different dtypes Float64 and Int64.

	. building graph for a.shape=(Float64) and b.shape=(Int64)

panic: Backend "xla": failed Add: dtype of first (Float64) and second (Int64) operands don't match

goroutine 1 [running]:
github.com/gomlx/gomlx/backends/xla.(*Builder).Add(0xc00028e000, {0xbd3ba0?, 0xc000290000?}, {0xbd3ba0, 0xc000290120})
	/home/janpf/Projects/gomlx/backends/xla/gen_standard_ops.go:32 +0x137
github.com/gomlx/gomlx/graph.Add(0xc0001ca580, 0xc0001ca630)
	/home/janpf/Projects/gomlx/graph/gen_backend_ops.go:176 +0x123
main.SumGraph(0xc0001ca580, 0xc0001ca630)
	 [[ Cell [4] Line 8 ]] /tmp/gonb_55ad4f33/main.go:17 +0x189
reflect.Value.call({0xbce4e0?, 0xc62b90?, 0xc0000d7850?}, {0xc4ba8a, 0x4}, {0xc000284150, 0x2, 0x0?})
	/snap/go/current/src/reflect/value.go:581 +0xca6
reflect.Value.Call({0xbce4e0?, 0xc62b90?, 0x0?}, {0xc000284150?, 0x0?, 0x0?})
	/snap/go/current/src/reflect/value.go:365 +0xb9
github.com/gomlx/gomlx/graph.(*Exec).createAndCacheGraph(0xc0001de120, {0xc0000da000, 0x2, 0xcb59a8?})
	/home/janpf/Projects/gomlx/graph/exec.go:536 +0x75c
github.com/gomlx/gomlx/graph.(*Exec).findOrCreateGraph(0xc0001de120, {0xc0000da000, 0x2, 0x2})
	/home/janpf/Projects/gomlx/graph/exec.go:595 +0x167
github.com/gomlx/gomlx/graph.(*Exec).compileAndExecute(0xc0001de120, 0x1, {0xc0000d7f20?, 0xc0000061c0?, 0xc0000d7f40?})
	/home/janpf/Projects/gomlx/graph/exec.go:436 +0x226
github.com/gomlx/gomlx/graph.(*Exec).CallWithGraph(...)
	/home/janpf/Projects/gomlx/graph/exec.go:386
github.com/gomlx/gomlx/graph.(*Exec).Call(0xcb59a8?, {0xc0000d7f20?, 0xbce4e0?, 0xc62b90?})
	/home/janpf/Projects/gomlx/graph/exec.go:372 +0x27
main.main()
	 [[ Cell [5] Line 3 ]] /tmp/gonb_55ad4f33/main.go:24 +0xbe
exit status 2

import "github.com/gomlx/exceptions"

err := exceptions.TryCatch[error](func() {_ = exec.Call(1.1, 2)})
if err != nil { … }

%%
onePlusExec := NewExec(backend, func (x *Node) *Node {
    return OnePlus(x) 
})
// exec.Call will return a *tensor.Tensor.
counter := onePlusExec.Call(0)[0]
// counter.Value() will first transfer counter to local with counter.Local().
fmt.Printf("counter.type=%T, counter.shape=%s, counter=%v\n", counter, counter.Shape(), counter.Value())
for ii := 0; ii < 10; ii++ {
    // Since the counter is not being used locally between the calls, the tensor will only use the 
    // device storage.
    counter = onePlusExec.Call(counter)[0]
}
// counter.Value() will first transfer the counter value locally, and then convert to a Go value.
fmt.Printf("counter=%v\n", counter.Value())

counter.type=*tensors.Tensor, counter.shape=(Int64), counter=1
counter=11

func f(x, y *Node) *Node {
    return Add(Square(x), Mul(x, y))
}

func gradOfF(x, y *Node) (output, gradX, gradY *Node) {
    output = f(x, y)
    reduced := ReduceAllSum(output) // In case x and y are not scalars.
    grads := Gradient(reduced, x, y)
    gradX, gradY = grads[0], grads[1] // df/dx, df/dy
    return output, gradX, gradY
}

%%
exec := NewExec(backend, gradOfF)
x := []float64{0, 1, 2}
y := []float64{10, 20, 30}
results := exec.Call(x, y)
fmt.Printf("f(x=%v, y=%v)=%v,\n\tdf/dx=%v,\n\tdf/dy=%v\n", x, y, results[0].Value(), results[1].Value(), results[2].Value())

f(x=[0 1 2], y=[10 20 30])=[0 21 64],
	df/dx=[10 22 34],
	df/dy=[0 1 2]

// Removing the previous definitions of `f` and `gradOfF`
%rm gradOfF f

. removed func gradOfF
. removed func f

import "flag"

var (
    flagA = flag.Float64("a", 1.0, "Value of a in the equation ax^2+bx+c")
    flagB = flag.Float64("b", 2.0, "Value of b in the equation ax^2+bx+c")
    flagC = flag.Float64("c", 4.0, "Value of c in the equation ax^2+bx+c")
    flagNumSteps = flag.Int("steps", 10, "Number of gradient descent steps to perform")
    flagLearningRate    = flag.Float64("lr", 0.1, "Initial learning rate.")
)

// f(x) = ax^2 + bx + c
func f(x *Node) *Node {
    f := MulScalar(Square(x), *flagA)
    f = Add(f, MulScalar(x, *flagB))
    f = AddScalar(f, *flagC)
    return f
}

// minimizeF does one gradient descent step on F by moving a variable "x",
// and returns the value of the function at the current "x".
func minimizeF(ctx *context.Context, graph *Graph) *Node {
    // Create or reuse existing variable "x" -- no graph operation is created with this, it's
    // only a reference.
    xVar := ctx.VariableWithShape("x", shapes.Make(dtypes.Float64))
    
    x := xVar.ValueGraph(graph)                        // Read variable for the current graph.
    y := f(x)                                          // Value of f(x).
    
    // Gradient always return a slice, we take the first element for grad of X.
    gradX := Gradient(y, x)[0] 
    
    // stepNum += 1
    stepNumVar := ctx.VariableWithValue("stepNum", 0.0)  // Creates the variable if not existing, or retrieve it if already exists.
    stepNum := stepNumVar.ValueGraph(graph)
    stepNum = OnePlus(stepNum)
    stepNumVar.SetValueGraph(stepNum)
    
    // step = -learningRate * gradX / Sqrt(stepNum)
    step := Div(gradX, Sqrt(stepNum))
    step = MulScalar(step, -*flagLearningRate)
    
    // x += step
    x = Add(x, step)
    xVar.SetValueGraph(x)
    return y  // f(x)
}

func Solve() {
    ctx := context.New()
    exec := context.NewExec(backend, ctx, minimizeF)
    for ii := 0; ii < *flagNumSteps-1; ii++ {
        _ = exec.Call()
    }
    y := exec.Call()[0]
    x := ctx.InspectVariable(ctx.Scope(), "x").Value()
    stepNum := ctx.InspectVariable(ctx.Scope(), "stepNum").Value()
    fmt.Printf("Minimum found at x=%g, f(x)=%g after %f steps.\n", x.Value(), y.Value(), stepNum.Value())
}

%% --a=1 --b=2 --c=3 --steps=10 --lr=0.5
Solve()

Minimum found at x=-0.9999999999999999, f(x)=2 after 10.000000 steps.

%% --a=2 --b=12 --c=20 --steps=10 --lr=0.5
Solve()

Minimum found at x=-3, f(x)=2 after 10.000000 steps.

const (
    CoefficientMu    = 0.0
    CoefficientSigma = 5.0
    BiasMu           = 1.0
    BiasSigma        = 10.0
)

import (
    "github.com/gomlx/gomlx/types/shapes"
    "github.com/gomlx/gomlx/types/tensors"
)

// initCoefficients chooses random coefficients and bias. These are the true values the model will
// attempt to learn.
func initCoefficients(backend backends.Backend, numVariables int) (coefficients, bias *tensors.Tensor) {
    e := NewExec(backend, func(g *Graph) (coefficients, bias *Node) {
        rngState := Const(g, RngState())
        rngState, coefficients = RandomNormal(rngState, shapes.Make(dtypes.Float64, numVariables))
        coefficients = AddScalar(MulScalar(coefficients, CoefficientSigma), CoefficientMu)
        rngState, bias = RandomNormal(rngState, shapes.Make(dtypes.Float64))
        bias = AddScalar(MulScalar(bias, BiasSigma), BiasMu)
        return
    })
    results := e.Call()
    coefficients, bias = results[0], results[1]
    return
}

%%
coef, bias := initCoefficients(backend, 3)
fmt.Printf("Example of target: coefficients=%0.3v, bias=%0.3v\n", coef.Value(), bias.Value())

Example of target: coefficients=[-3.4 -7.17 3.85], bias=-15.3

func buildExamples(backend backends.Backend, coef, bias *tensors.Tensor, numExamples int, noise float64) (inputs, labels *tensors.Tensor) {
    e := NewExec(backend, func(coef, bias *Node) (inputs, labels *Node) {
        g := coef.Graph()
        numFeatures := coef.Shape().Dimensions[0]

        // Random inputs (observations).
        rngState := Const(g, RngState())
        rngState, inputs = RandomNormal(rngState, shapes.Make(coef.DType(), numExamples, numFeatures))
        coef = ExpandDims(coef, 0)

        // Calculate perfect labels.
        labels = ReduceAndKeep(Mul(inputs, coef), ReduceSum, -1)
        labels = Add(labels, bias)
        if noise > 0 {
            // Add some noise to the labels.
            var noiseVector *Node
            rngState, noiseVector = RandomNormal(rngState, labels.Shape())
            noiseVector = MulScalar(noiseVector, noise)
            labels = Add(labels, noiseVector)
        }
        return
    })
    examples := e.Call(coef, bias)
    inputs, labels = examples[0], examples[1]
    return
}

%%
coef, bias := initCoefficients(backend, 3)
numExamples := 5
inputsTensor, labelsTensor := buildExamples(backend, coef, bias, numExamples, 0.2)
fmt.Printf("Target: coefficients=%0.3v, bias=%0.3v\n", coef.Value(), bias.Value()) 

fmt.Printf("%d dataset examples:\n", numExamples)
inputs := inputsTensor.Value().([][]float64)
labels := labelsTensor.Value().([][]float64)
for ii := 0; ii < numExamples; ii ++ {
    fmt.Printf("\tx=%0.3v; label=%0.3v\n", inputs[ii], labels[ii])
}

Target: coefficients=[7.07 1.15 -3.4], bias=4.38
5 dataset examples:
	x=[0.0184 1.35 -0.627]; label=[8.5]
	x=[0.876 -1.57 0.0963]; label=[8.32]
	x=[0.493 -0.332 0.89]; label=[3.87]
	x=[-0.535 0.923 1.01]; label=[-2.15]
	x=[0.701 0.357 0.696]; label=[7.72]

import "github.com/gomlx/gomlx/ml/train"

// TrivialDataset always returns the whole data.
type TrivialDataset struct {
    name string
    inputs, labels []*tensors.Tensor
}

var (
    // Assert Dataset implements train.Dataset.
    _ train.Dataset = &TrivialDataset{}
)
// Name implements train.Dataset.
func (ds *TrivialDataset) Name() string { return ds.name }

// Yield implements train.Dataset.
func (ds *TrivialDataset) Yield() (spec any, inputs, labels []*tensors.Tensor, err error) {
    return ds, ds.inputs, ds.labels, nil
}

// IsOwnershipTransferred tells the training loop that the dataset keeps ownership of the yielded tensors.
func (ds *TrivialDataset) IsOwnershipTransferred() bool {
    return false
}

// Reset implements train.Dataset.
func (ds *TrivialDataset) Reset() {}

type ModelFn func(ctx *context.Context, spec any, inputs []*graph.Node) (predictions []*graph.Node)

import "github.com/gomlx/gomlx/ml/context"

func modelGraph(ctx *context.Context, spec any, inputs []*Node) ([]*Node) {
    _ = spec  // Not needed here, we know the dataset.
    logits := layers.DenseWithBias(ctx, inputs[0], /* outputDim= */ 1)
    return []*Node{logits}
}

import (
    "os"
    "github.com/gomlx/gomlx/ui/commandline"
    "github.com/gomlx/gomlx/ml/layers/regularizers"
)

var (
    flagNumExamples  = flag.Int("num_examples", 10000, "Number of examples to generate")
    flagNumFeatures  = flag.Int("num_features", 3, "Number of features")
    flagNoise        = flag.Float64("noise", 0.2, "Noise in synthetic data generation")
    flagNumSteps     = flag.Int("steps", 100, "Number of gradient descent steps to perform")
    flagLearningRate    = flag.Float64("lr", 0.1, "Initial learning rate.")
)

// AttachToLoop decorators. It will be redefined later.
func AttachToLoop(loop *train.Loop) {
    commandline.AttachProgressBar(loop) // Attaches a progress bar to the loop. 
}

// TrainMain() does everything to train the linear model.
func TrainMain() {
    flag.Parse()

    // Select coefficients that we will try to predic.
    trueCoefficients, trueBias := initCoefficients(backend, *flagNumFeatures)
    fmt.Printf("Target: coefficients=%0.3v, bias=%0.3v\n", trueCoefficients.Value(), trueBias.Value()) 

    // Generate training data with noise.
    inputs, labels := buildExamples(backend, trueCoefficients, trueBias, *flagNumExamples, *flagNoise)
    fmt.Printf("Training data (inputs, labels): (%s, %s)\n\n", inputs.Shape(), labels.Shape())
    dataset := &TrivialDataset{"linear", []*tensors.Tensor{inputs}, []*tensors.Tensor{labels}}

    // Creates Context with learned weights and bias.
    ctx := context.New()
    ctx.SetParam(optimizers.ParamLearningRate, *flagLearningRate)  // = "learning_rate"
    ctx.SetParam(regularizers.ParamL2, 1e-3)  // 1e-3 of L2 regularization.

    // train.Trainer executes a training step.
    trainer := train.NewTrainer(backend, ctx, modelGraph,
        losses.MeanSquaredError,
        optimizers.StochasticGradientDescent(),
        nil, nil) // trainMetrics, evalMetrics
    loop := train.NewLoop(trainer)
    AttachToLoop(loop)

    // Loop for given number of steps. must.M1() panics, if loop.RunSteps returns an error.
    metrics := must.M1(loop.RunSteps(dataset, *flagNumSteps))
    _ = metrics  // We are not interested in them in this example.

    // Print learned coefficients and bias -- from the weights in the dense layer.
    fmt.Println()
    coefVar, biasVar := ctx.GetVariableByScopeAndName("/dense", "weights"), ctx.GetVariableByScopeAndName("/dense", "biases")
    learnedCoef, learnedBias := coefVar.Value(), biasVar.Value()
    fmt.Printf("Learned: coefficients=%0.3v, bias=%0.3v\n", learnedCoef.Value(), learnedBias.Value())
}

%%
TrainMain()

Target: coefficients=[-5.43 -1.65 -4.84], bias=6.74
Training data (inputs, labels): ((Float64)[10000 3], (Float64)[10000 1])

Training (100 steps):  100% [========================================] (4946 steps/s) [step=99] [loss+=0.144] [~loss+=5.162] [~loss=5.120]                       

Learned: coefficients=[[-5.31] [-1.62] [-4.73]], bias=[6.59]

import "github.com/gomlx/gomlx/ui/gonb/plotly"
import "os"

func AttachToLoop(loop *train.Loop) {
    commandline.AttachProgressBar(loop) // Attaches a progress bar to the loop. 
    _ = plotly.New().Dynamic().ScheduleExponential(loop, 50, 1.1)
}

%% --steps=5000
TrainMain()

Target: coefficients=[-1.74 3.39 1.23], bias=3.83
Training data (inputs, labels): ((Float64)[10000 3], (Float64)[10000 1])

Training (5000 steps):    3% [>.......................................] (5847 steps/s) [0s:0s] [step=164] [loss+=0.058] [~loss+=0.864] [~loss=0.851]

Training (5000 steps):  100% [========================================] (7731 steps/s) [step=4999] [loss+=0.056] [~loss+=0.056] [~loss=0.040]

Learned: coefficients=[[-1.73] [3.38] [1.23]], bias=[3.83]

batch_size := inputs[0].Shape().Dimensions[0]
    ...
    layer := Concatenate(allEmbeddings, -1)
    layer.AssertDims(batchSize, -1)  // 2D tensor, with batch size as the leading dimension.

GoMLX Tutorial¶

Note: Output Not Displaying in JupyterLab ?¶

Computation Graphs¶

Backends and executing Graphs with `Exec`¶

Error Handling¶

Tensors¶

Gradients¶

Variables and Context¶

Machine Learning (ML)¶

Dataset¶

ModelFn¶

Trainer and Loop¶

Training and Plotting¶

Other examples¶

Debugging¶

Good old "printf"¶

Errors reported with `panic` -- Go "exceptions".¶

Node Shape Asserts¶

Graph Execution Logging¶

Catching `NaN` and `Inf` in your training¶

More Debugging¶

GoMLX Tutorial¶

Note: Output Not Displaying in JupyterLab ?¶

Computation Graphs¶

Backends and executing Graphs with Exec¶

Error Handling¶

Tensors¶

Gradients¶

Variables and Context¶

Machine Learning (ML)¶

Dataset¶

ModelFn¶

Trainer and Loop¶

Training and Plotting¶

Other examples¶

Debugging¶

Good old "printf"¶

Errors reported with panic -- Go "exceptions".¶

Node Shape Asserts¶

Graph Execution Logging¶

Catching NaN and Inf in your training¶

More Debugging¶

Backends and executing Graphs with `Exec`¶

Errors reported with `panic` -- Go "exceptions".¶

Catching `NaN` and `Inf` in your training¶