122 lines
3.1 KiB
Go
122 lines
3.1 KiB
Go
|
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
|||
|
|
// SPDX-License-Identifier: MIT
|
|||
|
|
|
|||
|
|
package core
|
|||
|
|
|
|||
|
|
// TruncatedNormal returns an array of the given shape with values
|
|||
|
|
// drawn from a normal distribution truncated to ±2σ. Rejection sampling
|
|||
|
|
// redraws until a value lands inside the window. All draws share the
|
|||
|
|
// Generator for reproducibility; the dtype is float32. An invalid
|
|||
|
|
// shape yields a nil array, like the other no-error constructors, and
|
|||
|
|
// a std that is zero, negative or NaN is degenerate: the answer is the
|
|||
|
|
// all-zero array, the escape from a retry loop that negative and NaN
|
|||
|
|
// values could never leave.
|
|||
|
|
func TruncatedNormal(g *Generator, shape []int, mean, std float64) *Array {
|
|||
|
|
n, sh, err := checkedDims(shape)
|
|||
|
|
if err != nil {
|
|||
|
|
return nil
|
|||
|
|
}
|
|||
|
|
if !(std > 0) {
|
|||
|
|
// std <= 0 or NaN: the window is a point, inverted or undefined,
|
|||
|
|
// and the negative and NaN cases could never leave the retry loop
|
|||
|
|
// below, so the degenerate fallback returns the all-zero array.
|
|||
|
|
out := &Array{shape: sh, dt: Float32}
|
|||
|
|
out.alloc(n)
|
|||
|
|
return out
|
|||
|
|
}
|
|||
|
|
out := make([]float32, n)
|
|||
|
|
twoStd := 2.0 * std
|
|||
|
|
for i := range n {
|
|||
|
|
// The scalar draw consumes the generator exactly as one
|
|||
|
|
// Normal(g, 1, 0, std) call would (std is validated above, so
|
|||
|
|
// that call cannot error), without its per-element allocation.
|
|||
|
|
for {
|
|||
|
|
v := std * g.normalUnit()
|
|||
|
|
if v >= -twoStd && v <= twoStd {
|
|||
|
|
out[i] = float32(v + mean)
|
|||
|
|
break
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return &Array{shape: sh, dt: Float32, floats32: out}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// ZerosLike and OnesLike return a zero- or one-filled array with the
|
|||
|
|
// same shape and dtype as a. The result is a copy: mutations on it
|
|||
|
|
// do not affect a.
|
|||
|
|
func ZerosLike(a *Array) *Array {
|
|||
|
|
return fullLike(a, 0)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func OnesLike(a *Array) *Array {
|
|||
|
|
return fullLike(a, 1)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// fullLike creates a new array with a's shape and dtype, filled with v.
|
|||
|
|
// The narrow integer widths and bool take the same implicit-store cast
|
|||
|
|
// every other fill path carries: Go's conversion through int64, and
|
|||
|
|
// v != 0 for bool.
|
|||
|
|
func fullLike(a *Array, v float64) *Array {
|
|||
|
|
out := &Array{shape: a.Shape(), dt: a.dt}
|
|||
|
|
out.alloc(a.Len())
|
|||
|
|
switch a.dt {
|
|||
|
|
case Int:
|
|||
|
|
for i := range out.ints {
|
|||
|
|
out.ints[i] = int64(v)
|
|||
|
|
}
|
|||
|
|
case Bool:
|
|||
|
|
bv := v != 0
|
|||
|
|
for i := range out.bools {
|
|||
|
|
out.bools[i] = bv
|
|||
|
|
}
|
|||
|
|
case Int8:
|
|||
|
|
w := int8(int64(v))
|
|||
|
|
for i := range out.i8s {
|
|||
|
|
out.i8s[i] = w
|
|||
|
|
}
|
|||
|
|
case Uint8:
|
|||
|
|
w := uint8(int64(v))
|
|||
|
|
for i := range out.u8s {
|
|||
|
|
out.u8s[i] = w
|
|||
|
|
}
|
|||
|
|
case Int16:
|
|||
|
|
w := int16(int64(v))
|
|||
|
|
for i := range out.i16s {
|
|||
|
|
out.i16s[i] = w
|
|||
|
|
}
|
|||
|
|
case Uint16:
|
|||
|
|
w := uint16(int64(v))
|
|||
|
|
for i := range out.u16s {
|
|||
|
|
out.u16s[i] = w
|
|||
|
|
}
|
|||
|
|
case Int32:
|
|||
|
|
w := int32(int64(v))
|
|||
|
|
for i := range out.i32s {
|
|||
|
|
out.i32s[i] = w
|
|||
|
|
}
|
|||
|
|
case Uint32:
|
|||
|
|
w := uint32(int64(v))
|
|||
|
|
for i := range out.u32s {
|
|||
|
|
out.u32s[i] = w
|
|||
|
|
}
|
|||
|
|
case Float16:
|
|||
|
|
hv := HalfFromFloat64(v)
|
|||
|
|
for i := range out.halves {
|
|||
|
|
out.halves[i] = hv
|
|||
|
|
}
|
|||
|
|
case Float32:
|
|||
|
|
for i := range out.floats32 {
|
|||
|
|
out.floats32[i] = float32(v)
|
|||
|
|
}
|
|||
|
|
case Float:
|
|||
|
|
for i := range out.floats {
|
|||
|
|
out.floats[i] = v
|
|||
|
|
}
|
|||
|
|
default:
|
|||
|
|
for i := range out.complexes {
|
|||
|
|
out.complexes[i] = complex(v, 0)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return out
|
|||
|
|
}
|