Исправь полностью код, нейросеть сломана
```go
package neuralnetwork

import (
	"errors"
	"fmt"
	"math"
	"math/rand"
	"time"
)

// ActivationFunction определяет функцию активации и ее производную
type ActivationFunction struct {
	Function   func(float64) float64
	Derivative func(float64) float64
}

// tanhActivation и tanhDerivative реализуют функцию активации tanh и ее производную
func tanhActivation(x float64) float64 {
	return math.Tanh(x)
}

func tanhDerivative(x float64) float64 {
	return 1.0 - math.Pow(math.Tanh(x), 2)
}

// TanhActivation является экземпляром ActivationFunction для tanh
var TanhActivation = ActivationFunction{
	Function:   tanhActivation,
	Derivative: tanhDerivative,
}

// NeuralNetwork представляет простую рекуррентную нейронную сеть с одним скрытым слоем
type NeuralNetwork struct {
	activation          ActivationFunction
	inputSize           int
	hiddenSize          int
	outputSize          int
	hiddenStates        []float64
	weightsInputHidden  [][]float64
	weightsHiddenHidden [][]float64
	weightsHiddenOutput [][]float64
	trainRate           float64
}

// initializeWeights инициализирует матрицу весов с заданными размерами
func initializeWeights(rows, cols int) [][]float64 {
	weights := make([][]float64, rows)
	for i := 0; i < rows; i++ {
		weights[i] = make([]float64, cols)
		for j := 0; j < cols; j++ {
			weights[i][j] = rand.NormFloat64() * 0.1
		}
	}
	return weights
}

// updateWeights обновляет матрицу весов
func updateWeights(weights [][]float64, errors []float64, activations []float64, trainRate float64) error {
	for i := 0; i < len(weights); i++ {
		for j := 0; j < len(weights[i]); j++ {
			delta := trainRate * errors[j] * activations[i]
			if math.IsNaN(delta) {
				return fmt.Errorf("delta is NaN for weights[%d][%d]", i, j)
			}
			weights[i][j] += delta
		}
	}
	return nil
}

// NewNeuralNetwork создает новую нейронную сеть с инициализированными весами
func NewNeuralNetwork(
	activation ActivationFunction,
	inputSize, hiddenSize, outputSize int,
	trainRate float64,
) *NeuralNetwork {
	rand.Seed(time.Now().UnixNano())

	nn := &NeuralNetwork{
		activation:          activation,
		inputSize:           inputSize,
		hiddenSize:          hiddenSize,
		outputSize:          outputSize,
		hiddenStates:        make([]float64, hiddenSize),
		weightsInputHidden:  initializeWeights(inputSize, hiddenSize),
		weightsHiddenHidden: initializeWeights(hiddenSize, hiddenSize),
		weightsHiddenOutput: initializeWeights(hiddenSize, outputSize),
		trainRate:           trainRate,
	}

	return nn
}

// sampleFromDistribution выполняет выборку индекса из распределения вероятностей с учетом температуры
func sampleFromDistribution(output []float64, temperature float64) (int, error) {
	if temperature <= 0 {
		return -1, errors.New("temperature must be greater than 0")
	}

	// Применяем температуру
	logits := make([]float64, len(output))
	for i, v := range output {
		logits[i] = v / temperature
	}

	// Вычисляем softmax
	maxLogit := logits[0]
	for _, l := range logits {
		if l > maxLogit {
			maxLogit = l
		}
	}

	expSum := 0.0
	probs := make([]float64, len(logits))
	for i, l := range logits {
		probs[i] = math.Exp(l - maxLogit) // Вычитание maxLogit для численной стабильности
		expSum += probs[i]
	}

	for i := range probs {
		probs[i] /= expSum
	}

	// Генерируем случайное число
	randVal := rand.Float64()
	cumulative := 0.0
	for i, p := range probs {
		cumulative += p
		if randVal < cumulative {
			return i, nil
		}
	}

	// В случае округления возвращаем последний индекс
	return len(probs) - 1, nil
}

// Train обучает нейронную сеть на заданных данных
func (nn *NeuralNetwork) Train(epochs int, inputs [][]float64, targets [][]float64) error {
	if len(inputs) == 0 || len(targets) == 0 {
		return errors.New("inputs and targets must not be empty")
	}

	if len(inputs) != len(targets) {
		return errors.New("number of inputs and targets must match")
	}

	for i, input := range inputs {
		if len(input) != nn.inputSize {
			return fmt.Errorf("input at index %d has incorrect size", i)
		}
	}

	for i, target := range targets {
		if len(target) != nn.outputSize {
			return fmt.Errorf("target at index %d has incorrect size", i)
		}
	}

	for epoch := 0; epoch < epochs; epoch++ {
		fmt.Printf("Epoch: %d/%d [%.2f%%]\n", epoch+1, epochs, (float64(epoch+1)/float64(epochs))*100)
		for i := range inputs {
			input := inputs[i]
			target := targets[i]

			// Прямой проход (forward pass)
			hiddenInput := make([]float64, nn.hiddenSize)
			for h := 0; h < nn.hiddenSize; h++ {
				sum := 0.0
				for x := 0; x < nn.inputSize; x++ {
					sum += input[x] * nn.weightsInputHidden[x][h]
				}
				for hh := 0; hh < nn.hiddenSize; hh++ {
					sum += nn.hiddenStates[hh] * nn.weightsHiddenHidden[hh][h]
				}
				hiddenInput[h] = sum
				nn.hiddenStates[h] = nn.activation.Function(sum)
				if math.IsNaN(nn.hiddenStates[h]) {
					return fmt.Errorf("hidden state is NaN at epoch %d, input %d, hidden neuron %d", epoch, i, h)
				}
			}

			// Вычисление активаций выходного слоя
			outputs := make([]float64, nn.outputSize)
			for y := 0; y < nn.outputSize; y++ {
				sum := 0.0
				for h := 0; h < nn.hiddenSize; h++ {
					sum += nn.hiddenStates[h] * nn.weightsHiddenOutput[h][y]
				}
				outputs[y] = sum
				if math.IsNaN(outputs[y]) {
					return fmt.Errorf("output activation is NaN at epoch %d, input %d, output neuron %d", epoch, i, y)
				}
			}

			// Вычисление ошибок выходного слоя
			outputErrors := make([]float64, nn.outputSize)
			for y := 0; y < nn.outputSize; y++ {
				outputErrors[y] = target[y] - outputs[y]
			}

			// Обновление весов скрытого -> выходного слоя
			if err := updateWeights(nn.weightsHiddenOutput, outputErrors, nn.hiddenStates, nn.trainRate); err != nil {
				return err
			}

			// Вычисление ошибок скрытого слоя
			hiddenErrors := make([]float64, nn.hiddenSize)
			for h := 0; h < nn.hiddenSize; h++ {
				err := 0.0
				for y := 0; y < nn.outputSize; y++ {
					err += outputErrors[y] * nn.weightsHiddenOutput[h][y]
				}
				hiddenErrors[h] = err * nn.activation.Derivative(hiddenInput[h])
				if math.IsNaN(hiddenErrors[h]) {
					return fmt.Errorf("hidden error is NaN for hidden neuron %d", h)
				}
			}

			// Обновление рекуррентных весов скрытого слоя
			if err := updateWeights(nn.weightsHiddenHidden, hiddenErrors, nn.hiddenStates, nn.trainRate); err != nil {
				return err
			}

			// Обновление весов входного -> скрытого слоя
			if err := updateWeights(nn.weightsInputHidden, hiddenErrors, input, nn.trainRate); err != nil {
				return err
			}
		}
	}

	return nil
}

// Predict возвращает предсказание нейронной сети для заданных данных с выборкой
func (nn *NeuralNetwork) Predict(input []float64, temperature float64) ([]float64, error) {
	if len(input) != nn.inputSize {
		return nil, fmt.Errorf("input has incorrect size: expected %d, got %d", nn.inputSize, len(input))
	}

	// Сохраняем предыдущие состояния для восстановления после предсказания
	prevHiddenStates := make([]float64, nn.hiddenSize)
	copy(prevHiddenStates, nn.hiddenStates)

	// Прямой проход (forward pass)
	for h := 0; h < nn.hiddenSize; h++ {
		sum := 0.0
		for x := 0; x < nn.inputSize; x++ {
			sum += input[x] * nn.weightsInputHidden[x][h]
		}
		for hh := 0; hh < nn.hiddenSize; hh++ {
			sum += nn.hiddenStates[hh] * nn.weightsHiddenHidden[hh][h]
		}
		nn.hiddenStates[h] = nn.activation.Function(sum)
		if math.IsNaN(nn.hiddenStates[h]) {
			return nil, fmt.Errorf("hidden activation is NaN for hidden neuron %d", h)
		}
	}

	// Вычисление активаций выходного слоя
	outputs := make([]float64, nn.outputSize)
	for y := 0; y < nn.outputSize; y++ {
		sum := 0.0
		for h := 0; h < nn.hiddenSize; h++ {
			sum += nn.hiddenStates[h] * nn.weightsHiddenOutput[h][y]
		}
		outputs[y] = sum
		if math.IsNaN(outputs[y]) {
			return nil, fmt.Errorf("output activation is NaN for output neuron %d", y)
		}
	}

	// Выборку токена на основе распределения вероятностей с температурой
	tokenIndex, err := sampleFromDistribution(outputs, temperature)
	if err != nil {
		return nil, err
	}

	// Восстанавливаем предыдущие состояния
	copy(nn.hiddenStates, prevHiddenStates)

	// Создаем выходной вектор, где только выбранный токен имеет значение 1
	selectedOutput := make([]float64, nn.outputSize)
	selectedOutput[tokenIndex] = 1.0

	return selectedOutput, nil
}

// ResetHiddenStates сбрасывает состояния скрытого слоя
func (nn *NeuralNetwork) ResetHiddenStates() {
	for h := 0; h < nn.hiddenSize; h++ {
		nn.hiddenStates[h] = 0.0
	}
}

```