gradient.go

// Copyright 2019 The Gradient Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gradient

import (
	"math/cmplx"
	"math/rand"
)

// Copy copies src tensors into dst
func (context *Context[T]) Copy(k Continuation[T], node int, dst, src *V[T], options ...map[string]interface{}) bool {
	c := dst.CopyV(src)
	if k(c) {
		return true
	}
	for i, j := range c.D {
		src.D[i] += j
		dst.D[i] += j
	}
	return false
}

// Add adds two tensors
func (context *Context[T]) Add(k Continuation[T], node int, a, b *V[T], options ...map[string]interface{}) bool {
	length := len(b.X)
	c := a.Add(b)
	if k(c) {
		return true
	}

	for i, j := range c.D {
		a.D[i] += j
		b.D[i%length] += j
	}
	return false
}

// Sub subtracts two tensors
func (context *Context[T]) Sub(k Continuation[T], node int, a, b *V[T], options ...map[string]interface{}) bool {
	length := len(b.X)
	c := a.Sub(b)
	if k(c) {
		return true
	}
	for i, j := range c.D {
		a.D[i] += j
		b.D[i%length] -= j
	}
	return false
}

// Mul multiplies two tensors
func (context *Context[T]) Mul(k Continuation[T], node int, a, b *V[T], options ...map[string]interface{}) bool {
	width := a.S[0]
	sizeA, sizeB := len(a.X), len(b.X)
	c := a.MulT(b)
	if k(c) {
		return true
	}

	done := make(chan bool, 8)

	// a derivatives
	go func() {
		derivativeDone := make(chan bool, 8)
		derivatives := func(index int, ad []T) {
			rows, bi := a.S[1], 0
			for i := 0; i < sizeB; i += width {
				bv, cd := b.X[i:i+width], c.D[index+bi*rows]

				Axpy(cd, bv, ad)

				bi++
			}
			derivativeDone <- true
		}
		index := 0
		for j := 0; j < sizeA; j += width {
			ad := a.D[j : j+width]
			go derivatives(index, ad)
			index++
		}
		for j := 0; j < sizeA; j += width {
			<-derivativeDone
		}
		done <- true
	}()

	// b derivatives
	derivativeDone := make(chan bool, 8)
	derivatives := func(index int, bd []T) {
		for j := 0; j < sizeA; j += width {
			av, cd := a.X[j:j+width], c.D[index]

			Axpy(cd, av, bd)

			index++
		}
		derivativeDone <- true
	}
	index, rows := 0, a.S[1]
	for i := 0; i < sizeB; i += width {
		bd := b.D[i : i+width]
		go derivatives(index, bd)
		index += rows
	}
	for i := 0; i < sizeB; i += width {
		<-derivativeDone
	}
	<-done

	return false
}

// Square squares a tensor
func (context *Context[T]) Square(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	b := a
	width := a.S[0]
	sizeA, sizeB := len(a.X), len(b.X)
	c := a.SquareT()
	if k(c) {
		return true
	}

	// a derivatives
	{
		derivatives := func(index int, ad []T) {
			rows, bi := a.S[1], 0
			for i := 0; i < sizeB; i += width {
				bv, cd := b.X[i:i+width], c.D[index+bi*rows]

				Axpy(cd, bv, ad)

				bi++
			}
		}
		index := 0
		for j := 0; j < sizeA; j += width {
			ad := a.D[j : j+width]
			derivatives(index, ad)
			index++
		}
	}

	// b derivatives
	derivatives := func(index int, bd []T) {
		for j := 0; j < sizeA; j += width {
			av, cd := a.X[j:j+width], c.D[index]

			Axpy(cd, av, bd)

			index++
		}
	}
	index, rows := 0, a.S[1]
	for i := 0; i < sizeB; i += width {
		bd := b.D[i : i+width]
		derivatives(index, bd)
		index += rows
	}

	return false
}

// Euclidean computes the euclidean distance between all row vectors and all row vectors
func (context *Context[T]) Euclidean(k Continuation[T], node int, a, b *V[T], options ...map[string]interface{}) bool {
	width := a.S[0]
	sizeA, sizeB := len(a.X), len(b.X)
	c := a.Euclidean(b)
	if k(c) {
		return true
	}
	index := 0
	for i := 0; i < sizeA; i += width {
		for ii := 0; ii < sizeB; ii += width {
			av, bv, cx, ad, bd, d := a.X[i:i+width], b.X[ii:ii+width], c.X[index], a.D[i:i+width], b.D[ii:ii+width], c.D[index]
			for j, ax := range av {
				if cx == 0 {
					continue
				}
				ad[j] += (ax - bv[j]) * d / cx
				bd[j] += (bv[j] - ax) * d / cx
			}
			index++
		}
	}
	return false
}

// Hadamard computes the hadamard product of two tensors
func (context *Context[T]) Hadamard(k Continuation[T], node int, a, b *V[T], options ...map[string]interface{}) bool {
	length := len(b.X)
	c := a.Hadamard(b)
	if k(c) {
		return true
	}
	for i, j := range c.D {
		a.D[i] += j * b.X[i%length]
		b.D[i%length] += j * a.X[i]
	}
	return false
}

// T the transpose of the matrix
func (context *Context[T]) T(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	c := a.T()
	if k(c) {
		return true
	}
	i := 0
	for p := 0; p < a.S[0]; p++ {
		for q := 0; q < a.S[1]; q++ {
			a.D[q*a.S[0]+p] += c.D[i]
			i++
		}
	}
	return false
}

// H the conjugate transpose of the matrix
func (context *Context[T]) H(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	c := a.H()
	if k(c) {
		return true
	}
	i := 0
	for p := 0; p < a.S[0]; p++ {
		for q := 0; q < a.S[1]; q++ {
			switch ax := any(a.X[q*a.S[0]+p]).(type) {
			case float32:
				x := ax
				if x != 0 {
					x /= x
				}
				a.D[q*a.S[0]+p] += c.D[i] * any(x).(T)
			case float64:
				x := ax
				if x != 0 {
					x /= x
				}
				a.D[q*a.S[0]+p] += c.D[i] * any(x).(T)
			case complex64:
				ax = complex64(cmplx.Conj(complex128(ax)))
				x := real(ax)
				y := imag(ax)
				if x != 0 {
					x /= x
				}
				if y != 0 {
					y /= y
				}
				a.D[q*a.S[0]+p] += c.D[i] * any(complex(x, y)).(T)
			case complex128:
				ax = cmplx.Conj(ax)
				x := real(ax)
				y := imag(ax)
				if x != 0 {
					x /= x
				}
				if y != 0 {
					y /= y
				}
				a.D[q*a.S[0]+p] += c.D[i] * any(complex(x, y)).(T)
			}
			i++
		}
	}
	return false
}

// Slice a slice of the matrix
func (context *Context[T]) Slice(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	width := a.S[0]
	begin, end := *options[0]["begin"].(*int), *options[0]["end"].(*int)
	dd, ok := options[0]["d"].(*int)
	d := 0
	if ok {
		d = *dd
	}
	if d == 0 {
		c := a.Slice(begin, end, d)
		if k(c) {
			return true
		}
		index := 0
		ad := a.D[begin:end]
		for j := range ad {
			ad[j] += c.D[index]
			index++
		}
	} else if d == 2 {
		size := len(a.X)
		c := a.Slice(begin, end, d)
		if k(c) {
			return true
		}
		index := 0
		for i := 0; i < size; i += width {
			ad := a.D[i+begin : i+end]
			for j := range ad {
				ad[j] += c.D[index]
				index++
			}
		}
	}

	return false
}

// Concat concats two tensors
func (context *Context[T]) Concat(k Continuation[T], node int, a, b *V[T], options ...map[string]interface{}) bool {
	widthA, widthB := a.S[0], b.S[0]
	i, j := 0, 0
	c := a.Concat(b)
	if k(c) {
		return true
	}
	index, i, j := 0, 0, 0
	for r := 0; r < a.S[1]; r++ {
		ad, bd := a.D[i:i+widthA], b.D[j:j+widthB]
		for s := range ad {
			ad[s] = c.D[index]
			index++
		}
		for s := range bd {
			bd[s] = c.D[index]
			index++
		}
		i += widthA
		j += widthB
	}
	return false
}

// Dropout is a dropout regularization function
func (context *Context[T]) Dropout(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	size, width := len(a.X), a.S[0]
	rng := options[0]["rng"].(*rand.Rand)
	drop := .1
	if options[0]["drop"] != nil {
		drop = *options[0]["drop"].(*float64)
	}
	drops := make([]int, width)
	for i := range drops {
		if rng.Float64() > drop {
			drops[i] = 1
		}
	}
	c := a.Dropout(drop, drops)
	if k(c) {
		return true
	}
	for i := 0; i < size; i += width {
		for j := range a.D[i : i+width] {
			if drops[j] == 1 {
				a.D[i+j] += c.D[i+j]
			}
		}
	}
	return false
}

// DropoutMatrix is a dropout regularization function
func (context *Context[T]) DropoutMatrix(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	size, width := len(a.X), a.S[0]
	rng := options[0]["rng"].(*rand.Rand)
	drop := .1
	if options[0]["drop"] != nil {
		drop = *options[0]["drop"].(*float64)
	}
	drops := make([]int, size)
	for i := range drops {
		if rng.Float64() > drop {
			drops[i] = 1
		}
	}
	c := a.DropoutMatrix(drop, drops)
	if k(c) {
		return true
	}
	for i := 0; i < size; i += width {
		for j := range a.D[i : i+width] {
			if drops[i+j] == 1 {
				a.D[i+j] += c.D[i+j]
			}
		}
	}
	return false
}

// Sin the sine of a number
func (context *Context[T]) Sin(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	c := a.Sin()
	if k(c) {
		return true
	}
	for i, j := range c.D {
		a.D[i] += j * Cos(a.X[i])
	}
	return false
}

// Cos the cosine of a tensor
func (context *Context[T]) Cos(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	c := a.Cos()
	if k(c) {
		return true
	}
	for i, j := range c.D {
		a.D[i] -= j * Sin(a.X[i])
	}
	return false
}

// Exp the base e exponential of a tensor
func (context *Context[T]) Exp(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	c := a.Exp()
	if k(c) {
		return true
	}
	for i, j := range c.D {
		a.D[i] += j * c.X[i]
	}
	return false
}

// Log the natural logarithm of a tensor
func (context *Context[T]) Log(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	c := a.Log()
	if k(c) {
		return true
	}
	for i, j := range c.D {
		a.D[i] += j / a.X[i]
	}
	return false
}

// Sqrt is the sqrt of a number
func (context *Context[T]) Sqrt(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	c := a.Sqrt()
	if k(c) {
		return true
	}
	for i, j := range c.D {
		a.D[i] += j / (2 * c.X[i])
	}
	return false
}

// Inv is the inverse of a number
func (context *Context[T]) Inv(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	c := a.Inv()
	if k(c) {
		return true
	}
	for i, j := range c.D {
		if a.X[i] == 0 {
			continue
		}
		a.D[i] += -j / (a.X[i] * a.X[i])
	}
	return false
}

// Sigmoid computes the sigmoid of a vector
func (context *Context[T]) Sigmoid(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	c := a.Sigmoid()
	if k(c) {
		return true
	}
	for i, j := range c.D {
		cx := c.X[i]
		a.D[i] += j * cx * (1 - cx)
	}
	return false
}

// TanH the hyperbolic tangent of a tensor
func (context *Context[T]) TanH(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	c := a.TanH()
	if k(c) {
		return true
	}
	for i, j := range c.D {
		cx := c.X[i]
		a.D[i] += j * (1 - cx*cx)
	}
	return false
}

// Softplus the softplus activation function
func (context *Context[T]) Softplus(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	c := a.Softplus()
	if k(c) {
		return true
	}
	for i, j := range c.D {
		a.D[i] += j / (1 + Exp(-a.X[i]))
	}
	return false
}

// Everett computes the split reality activation function
func (context *Context[T]) Everett(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	c := a.Everett()
	if k(c) {
		return true
	}

	for i, j := range c.D {
		if c.X[i] != 0 || (c.X[i&^1] == 0 && c.X[i|1] == 0) {
			a.D[i>>1] += j
		}
	}
	return false
}

// EverettReLu computes an adapter relu
func (context *Context[T]) EverettReLu(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	c := a.EverettReLu()
	if k(c) {
		return true
	}
	for i, j := range c.D {
		if c.X[i] != 0 {
			a.D[i>>1] += j
		}
	}
	return false
}

// ReLu computes the rectified linear activation function
func (context *Context[T]) ReLu(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	c := a.ReLu()
	if k(c) {
		return true
	}
	for i, j := range c.D {
		if c.X[i] != 0 {
			a.D[i] += j
		}
	}
	return false
}

const (
	// S is the scaling factor for the softmax
	S = 1.0 - 1e-300
)

// Softmax is the softmax function for big numbers
func (context *Context[T]) Softmax(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	S := S
	if len(options) > 0 {
		s, ok := options[0]["S"]
		if ok {
			S = s.(float64)
		}
	}
	c := a.Softmax(S)
	if k(c) {
		return true
	}
	for i, d := range c.D {
		cx := c.X[i]
		for j := range c.X {
			if j == i {
				a.D[j] += d * cx * (1 - cx)
			} else {
				a.D[j] -= d * cx * c.X[j]
			}
		}
	}
	return false
}

// Sum sums a vector
func (context *Context[T]) Sum(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	c := a.Sum()
	if k(c) {
		return true
	}
	d := c.D[0]
	for i := range a.D {
		a.D[i] += d
	}
	return false
}

// SumRows sums the rows of the matrix
func (context *Context[T]) SumRows(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	size, width := len(a.X), a.S[0]
	c := a.SumRows()
	if k(c) {
		return true
	}
	for i := 0; i < size; i += width {
		for j := range a.D[i : i+width] {
			a.D[i+j] += c.D[j]
		}
	}
	return false
}

// Quadratic computes the quadratic cost of two tensors
func (context *Context[T]) Quadratic(k Continuation[T], node int, a, b *V[T], options ...map[string]interface{}) bool {
	width := a.S[0]
	size := len(a.X)
	c := a.Quadratic(b)
	if k(c) {
		return true
	}
	index := 0
	for i := 0; i < size; i += width {
		av, bv, ad, bd, d := a.X[i:i+width], b.X[i:i+width], a.D[i:i+width], b.D[i:i+width], c.D[index]
		for j, ax := range av {
			ad[j] += (ax - bv[j]) * d
			bd[j] += (bv[j] - ax) * d
		}
		index++
	}
	return false
}

// CrossEntropy computes the cross entropy cost of two tensors
func (context *Context[T]) CrossEntropy(k Continuation[T], node int, a, b *V[T], options ...map[string]interface{}) bool {
	width := a.S[0]
	size := len(a.X)
	c := a.CrossEntropy(b)
	if k(c) {
		return true
	}
	index := 0
	for i := 0; i < size; i += width {
		av, bv, ad, bd, d := a.X[i:i+width], b.X[i:i+width], a.D[i:i+width], b.D[i:i+width], c.D[index]
		for j, ax := range av {
			bx := bv[j]
			if bx == 1 {
				ad[j] -= d / (ax + .001)
				bd[j] -= Log(ax+.001) * d
			} else {
				ad[j] += d / (1 - ax + .001)
				bd[j] -= Log(1-ax+.001) * d
			}
		}
		index++
	}
	return false
}

// Similarity computes the cosine similarity cost of two tensors
func (context *Context[T]) Similarity(k Continuation[T], node int, a, b *V[T], options ...map[string]interface{}) bool {
	width := a.S[0]
	size := len(a.X)
	c, ab, aa, bb := a.Similarity(b)
	if k(c) {
		return true
	}
	index := 0
	for i := 0; i < size; i += width {
		av, bv, ad, bd, cd := a.X[i:i+width], b.X[i:i+width], a.D[i:i+width], b.D[i:i+width], c.D[index]
		sumAB, sumAA, sumBB := ab[index], aa[index], bb[index]
		denominator := Sqrt(sumAA) * Sqrt(sumBB)
		for j, ax := range av {
			bx := bv[j]
			ad[j] += cd * (bx/denominator - ax*sumAB/(sumAA*denominator))
			bd[j] += cd * (ax/denominator - bx*sumAB/(sumBB*denominator))
		}
		index++
	}
	return false
}

// Orthogonality computes the cosine similarity between all vectors
func (context *Context[T]) Orthogonality(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	size, width := len(a.X), a.S[0]
	c, ab, aa, bb := a.Orthogonality()
	if k(c) {
		return true
	}
	index := 0
	for i := 0; i < size; i += width {
		for j := i + width; j < size; j += width {
			cd, sumAB, sumAA, sumBB := c.D[index], ab[index], aa[index], bb[index]
			denominator := Sqrt(sumAA) * Sqrt(sumBB)
			for k := 0; k < width; k++ {
				ax, bx := a.X[i+k], a.X[j+k]
				a.D[i+k] += cd * (bx/denominator - ax*sumAB/(sumAA*denominator))
				a.D[j+k] += cd * (ax/denominator - bx*sumAB/(sumBB*denominator))
			}
			index++
		}
	}
	return false
}

// Entropy computes the entropy of the vectors
func (context *Context[T]) Entropy(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	size, width := len(a.X), a.S[0]
	c := a.Entropy()
	if k(c) {
		return true
	}
	index := 0
	for i := 0; i < size; i += width {
		cd := c.D[index]
		for k := 0; k < width; k++ {
			ax := a.X[i+k]
			a.D[i+k] -= cd * (Log(ax) + 1)
		}
		index++
	}
	return false
}

// Variance computes the variance of the vectors
func (context *Context[T]) Variance(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	size, width := len(a.X), a.S[0]
	c, means := a.Variance()
	n := Convert[T](float64(width))
	if k(c) {
		return true
	}
	index, nn := 0, n*n
	for i := 0; i < size; i += width {
		cd, mean := c.D[index], means[index]
		for j := 0; j < width; j++ {
			sum := T(0.0)
			for k := 0; k < width; k++ {
				d := a.X[i+k] - mean
				if j == k {
					d *= (n - 1)
				} else {
					d *= -1
				}
				sum += d
			}
			a.D[i+j] += cd * 2 * sum / nn
		}
		index++
	}
	return false
}

// Abs computes the absolute value of the tensor
func (context *Context[T]) Abs(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	c := a.Abs()
	if k(c) {
		return true
	}
	for i, cd := range c.D {
		sign := Convert[T](float64(Sign(a.X[i])))
		a.D[i] += cd * sign
	}
	return false
}

// Quantize quantizes the values
func (context *Context[T]) Quant(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	c := a.Quant(context)
	if k(c) {
		return true
	}
	for i, cd := range c.D {
		a.D[i] += cd
	}
	return false
}

// Avg computes the average of the tensor
func (context *Context[T]) Avg(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	total := Convert[T](float64(len(a.X)))
	c := a.Avg()
	if k(c) {
		return true
	}
	d := c.D[0] / total
	for i := range a.D {
		a.D[i] += d
	}
	return false
}

// Combines two complex tensors to a complex tensor
func (context *Context[T]) Complex(k Continuation[T], node int, a, b *V[T], options ...map[string]interface{}) bool {
	length := len(b.X)
	c := a.Complex(b)
	if k(c) {
		return true
	}
	for i, cD := range c.D {
		switch any(cD).(type) {
		case complex64:
			aX, bX := cmplx.Abs(complex128(any(a.X[i]).(complex64))), cmplx.Phase(complex128(any(b.X[i%length]).(complex64)))
			cos, sin := Cos(bX), Sin(bX)
			a.D[i] += cD * Convert[T](cos+sin)
			b.D[i%length] += cD * Convert[T](aX*(cos-sin))
		case complex128:
			aX, bX := cmplx.Abs(any(a.X[i]).(complex128)), cmplx.Phase(any(b.X[i%length]).(complex128))
			cos, sin := Cos(bX), Sin(bX)
			a.D[i] += cD * Convert[T](cos+sin)
			b.D[i%length] += cD * Convert[T](aX*(cos-sin))
		}
	}
	return false
}

// Phase computes the phase of a complex tensor
func (context *Context[T]) Phase(k Continuation[T], node int, a *V[T], options ...map[string]interface{}) bool {
	c := a.Phase()
	if k(c) {
		return true
	}
	for i, cD := range c.D {
		ax := a.X[i]
		switch ax := any(ax).(type) {
		case complex64:
			aa := real(ax)
			bb := imag(ax)
			cc := (aa*aa + bb*bb)
			a.D[i] += cD * any(complex((bb*bb-aa*aa)/(cc*cc), 0)).(T)
		case complex128:
			aa := real(ax)
			bb := imag(ax)
			cc := (aa*aa + bb*bb)
			a.D[i] += cD * any(complex((bb*bb-aa*aa)/(cc*cc), 0)).(T)
		}

	}
	return false
}