<commit_before>package bwmf

import (
	"fmt"
	// "math"
	"testing"

	pb "github.com/taskgraph/taskgraph/example/bwmf/proto"
	"github.com/taskgraph/taskgraph/op"
)

func TestEvaluation(t *testing.T) {

	h := op.NewVecParameter(4)
	g := h.CloneWithoutCopy()
	h.Set(0, 1.0)
	h.Set(1, 1.0)
	h.Set(2, 1.0)
	h.Set(3, 1.0)

	kld_loss := newKLDivLoss()

	loss_val := kld_loss.Evaluate(h, g)

	fmt.Println("loss val is ", loss_val)
	fmt.Println("gradient is ", g)

	/*
		if math.Abs(float64(loss_val)-4.6) > 5e-2 {
			t.Errorf("Loss value incorrect: actual %f, expected 6.06543", loss_val)
		}

		if math.Abs(float64(g.Get(0))+1.01) > 1e-2 || math.Abs(float64(g.Get(1))+1.01) > 1e-2 {
			t.Errorf("Gradient incorrect: actual {%f, %f}, expected { -1.01, -1.01 }", g.Get(0), g.Get(1))
		}
	*/

	h.Set(0, 0.0)
	h.Set(1, 1.0)
	h.Set(2, 1.0)
	h.Set(3, 0.0)

	loss_val = kld_loss.Evaluate(h, g)

	fmt.Println("loss val is ", loss_val)
	fmt.Println("gradient is ", g)

}

func newKLDivLoss() *KLDivLoss {
	m := uint32(3)
	n := uint32(2)
	k := uint32(2)

	// v is
	// -------------
	// | 0.0 | 1.0 |
	// |-----|-----|
	// | 1.0 | 0.0 |
	// |-----|-----|
	// | 0.5 | 0.5 |
	// |-----|-----|
	v := &pb.MatrixShard{
		IsSparse: true,
		M: 2,
		N: 3,
		Val: []float32 {1.0, 1.0, 0.5, 0.5},
		Ir: []uint32 {1, 0, 0, 1},
		Jc: []uint32 {0, 1, 2, 4},
	}

	// w is composed by two shards:
	//
	// shard 0:
	//  -------------
	//  | 1.0 | 0.0 |
	//  -------------
	//  | 0.0 | 1.0 |
	//  -------------
	// and shard 1:
	//  -------------
	//  | 0.5 | 0.5 |
	//  -------------
	//
	w := []*pb.MatrixShard{
		&pb.MatrixShard{
			M: 2,
			N: n,
			Val: []float32 {1.0, 0, 0.0, 1.0},
		},
		&pb.MatrixShard{
			M: 1,
			N: n,
			Val: []float32 {0.5, 0.5},
		},
	}

	return NewKLDivLoss(v, w, m, n, k, 1e-6)
}<commit_msg>bwmf: check the result of kldiv in unit tests<commit_after>package bwmf

import (
	"fmt"
	"math"
	"testing"

	pb "github.com/taskgraph/taskgraph/example/bwmf/proto"
	"github.com/taskgraph/taskgraph/op"
)

func TestEvaluation(t *testing.T) {

	h := op.NewVecParameter(4)
	g := h.CloneWithoutCopy()
	h.Set(0, 1.0)
	h.Set(1, 1.0)
	h.Set(2, 1.0)
	h.Set(3, 1.0)

	kld_loss := newKLDivLoss()

	loss_val := kld_loss.Evaluate(h, g)

	fmt.Println("loss val is ", loss_val)
	fmt.Println("gradient is ", g)

	expectedLoss := 5.999997
	if math.Abs(float64(loss_val)-expectedLoss) > 1e-5 {
		t.Errorf("Loss value incorrect: actual %f, expected %f", loss_val, expectedLoss)
	}

	expectedGrad := []float32 { 1.2499998, 0.24999973, 0.24999973, 1.2499998 }
	if euclideanDist(g.Data(), expectedGrad) > 1e-5 {
		t.Errorf("Gradient incorrect: actual %v, expected %v", g.Data(), expectedGrad)
	}

	h.Set(0, 0.0)
	h.Set(1, 1.0)
	h.Set(2, 1.0)
	h.Set(3, 0.0)

	loss_val = kld_loss.Evaluate(h, g)

	fmt.Println("loss val is ", loss_val)
	fmt.Println("gradient is ", g)

	expectedLoss = 3.6931434
	if math.Abs(float64(loss_val)-expectedLoss) > 1e-5 {
		t.Errorf("Loss value incorrect: actual %f, expected %f", loss_val, expectedLoss)
	}

	expectedGrad = []float32 { 1, 0, 0, 1 }
	if euclideanDist(g.Data(), expectedGrad) > 1e-5 {
		t.Errorf("Gradient incorrect: actual %v, expected %v", g.Data(), expectedGrad)
	}
}

func newKLDivLoss() *KLDivLoss {
	m := uint32(3)
	n := uint32(2)
	k := uint32(2)

	// v is
	// -------------
	// | 0.0 | 1.0 |
	// |-----|-----|
	// | 1.0 | 0.0 |
	// |-----|-----|
	// | 0.5 | 0.5 |
	// |-----|-----|
	v := &pb.MatrixShard{
		IsSparse: true,
		M: 2,
		N: 3,
		Val: []float32 {1.0, 1.0, 0.5, 0.5},
		Ir: []uint32 {1, 0, 0, 1},
		Jc: []uint32 {0, 1, 2, 4},
	}

	// w is composed by two shards:
	//
	// shard 0:
	//  -------------
	//  | 1.0 | 0.0 |
	//  -------------
	//  | 0.0 | 1.0 |
	//  -------------
	// and shard 1:
	//  -------------
	//  | 0.5 | 0.5 |
	//  -------------
	//
	w := []*pb.MatrixShard{
		&pb.MatrixShard{
			M: 2,
			N: n,
			Val: []float32 {1.0, 0.0, 0.0, 1.0},
		},
		&pb.MatrixShard{
			M: 1,
			N: n,
			Val: []float32 {0.5, 0.5},
		},
	}

	return NewKLDivLoss(v, w, m, n, k, 1e-6)
}

func euclideanDist(x, y []float32) float64 {
	if len(x) != len(y) {
		return math.Inf(1)
	}

	val := 0.0

	for i, xi := range x {
		val += float64((xi-y[i])*(xi-y[i]))
	}
	return val
}
{'fraction_non_alphanumeric': 0.2612118320610687, 'fraction_numerical': 0.07657442748091603, 'mean_word_length': 2.122114668652271, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 10, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '794809', 'n_tokens_mistral': 2265, 'n_tokens_neox': 1982, 'n_words': 405}