TensorFrost 0.7.0.dev5__cp313-cp313-win_amd64.whl

TensorFrost/__init__.py

@@ -0,0 +1,8 @@
+from .TensorFrost import *
+
+from . import optimizers
+from . import regularizers
+from . import clipping
+from . import random
+from . import sort
+from .default import *

TensorFrost/clipping.py

@@ -0,0 +1,5 @@
+from .optimizers import *
+
+clamp = ModuleOptimizer.ClippingType.Clamp
+norm = ModuleOptimizer.ClippingType.Norm
+none = ModuleOptimizer.ClippingType.None_

TensorFrost/default.py

@@ -0,0 +1,14 @@
+from . import TensorFrost as tf
+
+def zeros_like(tensor):
+    return tf.zeros(tensor.shape, tensor.type)
+
+def eye(n):
+    i, j = tf.indices([n, n])
+    return tf.select(i == j, 1.0, 0.0)
+
+def eye_like(tensor):
+    return eye(tensor.shape[0])
+
+def ones_like(tensor):
+    return tf.ones(tensor.shape, tensor.type)

TensorFrost/optimizers.py

@@ -0,0 +1,219 @@
+from . import TensorFrost as tf
+
+class ModuleOptimizer(tf.Module):
+    class OptimizerType:
+        ADAM = 0
+        SGD = 1
+        RMSProp = 2
+
+    class RegularizerType:
+        None_ = 0
+        L1 = 1
+        L2 = 2
+
+    class ClippingType:
+        Clamp = 0
+        Norm = 1
+        None_ = 2
+
+    def __init__(self, optimizer_type, regularizer_type, net, params):
+        super().__init__()
+        self.optimizer_type = optimizer_type
+        self.regularizer_type = regularizer_type
+        self.clipping_type = self.ClippingType.Clamp
+        self.epsilon = 1e-8
+
+        # Set passed parameters as attributes
+        self.net = net
+        for k, v in params.items():
+            setattr(self, k, v)
+
+        # Initialize t
+        t = tf.Parameter([1], tf.float32, False)  # mimic Parameter({1}, TFType::Float, false)
+        self.t = t
+
+        self.initializeOptimizer(net)
+
+    def set_clipping_type(self, ctype):
+        self.clipping_type = ctype
+
+    def initializeOptimizer(self, net):
+        net_params = net.parameters()
+        requires_grads = net.requires_grads_list()
+
+        if self.optimizer_type == self.OptimizerType.ADAM:
+            self.initializeParameterArray("m", net_params, requires_grads)
+            self.initializeParameterArray("v", net_params, requires_grads)
+        elif self.optimizer_type == self.OptimizerType.SGD:
+            # No additional parameters needed
+            pass
+        elif self.optimizer_type == self.OptimizerType.RMSProp:
+            self.initializeParameterArray("v", net_params, requires_grads)
+
+    def initializeParameterArray(self, name, net_params, requires_grads):
+        arr = tf.ParameterArray()
+
+        for i, param in enumerate(net_params):
+            if not requires_grads[i]:
+                continue
+
+            new_param = tf.Parameter(param.shape, tf.float32, False)
+            arr[i] = new_param
+
+        setattr(self, name, arr)
+
+    def assert_parameters(self):
+        net_params = self.net.parameters()
+        requires_grads = self.net.requires_grads_list()
+        self.assertParameterArray("m", net_params, requires_grads)
+        self.assertParameterArray("v", net_params, requires_grads)
+
+    def gradient_norm(self, grad):
+        # sum of squares
+        g = grad * grad
+        shape = grad.shape
+        num_dims = len(shape)
+        for i in range(num_dims):
+            g = tf.sum(g)
+        return tf.sqrt(g)
+
+    def assertParameterArray(self, name, net_params, requires_grads):
+        if hasattr(self, name):
+            arr = getattr(self, name)
+            for i, param in enumerate(net_params):
+                if not requires_grads[i]:
+                    continue
+                arr_item = arr[i]
+                arr_item = tf.assert_tensor(arr_item, param.shape, param.type)
+                arr[i] = arr_item
+
+    def step(self, *args):
+        # Overloaded step:
+        # step(X, Y) or step(loss)
+        if len(args) == 2:
+            X, Y = args
+            loss = self.net.loss(X, Y)
+            self._step(loss)
+            return loss
+        elif len(args) == 1:
+            (loss,) = args
+            self._step(loss)
+        else:
+            raise ValueError("Invalid arguments to step")
+
+    def _step(self, loss):
+        # Increment t by 1
+        self.t = self.t + 1.0
+
+        net = self.net
+        net_params = net.parameters()
+        requires_grads = net.requires_grads_list()
+
+        learning_rate = self.learning_rate
+        grad_clip = self.grad_clip
+        has_clip = isinstance(grad_clip, float) and grad_clip > 0.0
+
+        for i, param in enumerate(net_params):
+            if not requires_grads[i]:
+                continue
+
+            grad = tf.grad(loss, param)
+            if has_clip:
+                if self.clipping_type == self.ClippingType.Clamp:
+                    grad = tf.clamp(grad, -grad_clip, grad_clip)
+                elif self.clipping_type == self.ClippingType.Norm:
+                    grad_norm = tf.max(1e-6, self.gradient_norm(grad))
+                    grad = grad * tf.min(1.0, grad_clip / grad_norm)
+
+            if self.optimizer_type == self.OptimizerType.ADAM:
+                update = self.adam_update(i, param, grad, self.t, learning_rate)
+            elif self.optimizer_type == self.OptimizerType.SGD:
+                update = self.sgd_update(param, grad, learning_rate)
+            elif self.optimizer_type == self.OptimizerType.RMSProp:
+                update = self.rmsprop_update(i, param, grad, learning_rate)
+            else:
+                raise RuntimeError("Unknown optimizer type")
+
+            # Apply regularization if needed
+            if self.regularizer_type == self.RegularizerType.L1:
+                param = param - learning_rate * self.reg * tf.sign(param)
+            elif self.regularizer_type == self.RegularizerType.L2:
+                param = param - learning_rate * self.reg * param
+
+            # Update parameter with computed update
+            param = param - update
+            net_params[i] = param
+
+        net.update_parameters(net_params)
+
+    def adam_update(self, i, param, grad, t, learning_rate):
+        beta1 = tf.float(self.beta1)
+        beta2 = tf.float(self.beta2)
+
+        m = self.m[i]
+        v = self.v[i]
+
+        m = tf.lerp(grad, m, beta1)
+        v = tf.lerp(grad * grad, v, beta2)
+
+        # t is a Parameter with shape [1]; get the scalar
+        t_val = self.t[0]
+        mhat = m / (1.0 - tf.pow(beta1, t_val))
+        vhat = v / (1.0 - tf.pow(beta2, t_val))
+
+        self.m[i] = m
+        self.v[i] = v
+
+        return learning_rate * mhat / (tf.sqrt(vhat) + self.epsilon)
+
+    def sgd_update(self, param, grad, learning_rate):
+        return learning_rate * grad
+
+    def rmsprop_update(self, i, param, grad, learning_rate):
+        decay = tf.float(self.decay)
+
+        v = self.v[i]
+        v = tf.lerp(grad * grad, v, decay)
+        self.v[i] = v
+
+        return (grad * learning_rate) / (tf.sqrt(v) + self.epsilon)
+
+
+def adam(net, reg_type=ModuleOptimizer.RegularizerType.None_, learning_rate=0.001, beta1=0.9, beta2=0.999, clip=0.0, reg=0.0):
+    return ModuleOptimizer(
+        ModuleOptimizer.OptimizerType.ADAM,
+        reg_type,
+        net,
+        {
+            "learning_rate": learning_rate,
+            "beta1": beta1,
+            "beta2": beta2,
+            "grad_clip": clip,
+            "reg": reg,
+        }
+    )
+
+def sgd(net, reg_type=ModuleOptimizer.RegularizerType.None_, learning_rate=0.001, clip=0.0, reg=0.0):
+    return ModuleOptimizer(
+        ModuleOptimizer.OptimizerType.SGD,
+        reg_type,
+        net,
+        {
+            "learning_rate": learning_rate,
+            "grad_clip": clip,
+            "reg": reg,
+        }
+    )
+
+def rmsprop(net, reg_type=ModuleOptimizer.RegularizerType.None_, learning_rate=0.001, decay=0.9, clip=0.0, reg=0.0):
+    return ModuleOptimizer(
+        ModuleOptimizer.OptimizerType.RMSProp,
+        reg_type,
+        net,
+        {
+            "learning_rate": learning_rate,
+            "decay": decay,
+            "grad_clip": clip,
+            "reg": reg,
+        }
+    )

TensorFrost/random.py

@@ -0,0 +1,45 @@
+from . import TensorFrost as tf
+
+def randn2(shape, seed=0):
+    #Box-Muller transform
+    r1 = tf.random_value(shape, seed=seed)
+    r2 = tf.random_value(shape, seed=tf.hash(seed))
+    rho = tf.sqrt(-2.0*tf.log(tf.max(1e-6, r1)))
+    theta = 2.0*tf.pi*r2
+    return rho*tf.cos(theta), rho*tf.sin(theta)
+
+def randn(shape, seed=0):
+    return randn2(shape, seed=seed)[0]
+
+def rand(shape, seed=0):
+    return tf.random_value(shape, seed=seed)
+
+def randn_like(tensor, seed=0):
+    return randn(tensor.shape, seed=seed)
+
+def rand_like(tensor, seed=0):
+    return rand(tensor.shape, seed=seed)
+
+def rand_int(seed, max_value):
+    return tf.int(tf.pcg(tf.uint(seed)) % tf.uint(max_value))
+
+def xor_swap(idx, n, seed):
+    xor_seed = rand_int(seed, n)
+    xor_idx = (idx ^ xor_seed)
+    max_idx = tf.max(idx, xor_idx)
+    min_idx = tf.min(idx, xor_idx)
+    swap = rand_int(min_idx * 451 + seed, 2) == 0
+    return tf.select(swap & (max_idx < n), xor_idx, idx)
+
+def reverse(idx, n):
+    return n - 1 - idx
+
+def shuffle(idx, n, seed = 0, iters = 16):
+    for i in range(iters):
+        idx = xor_swap(idx, n, seed + i)
+        idx = reverse(idx, n)
+    return idx
+
+def permutation(n, seed = 0):
+    idx = tf.indices([n])[0]
+    return shuffle(idx, n, seed)

TensorFrost/regularizers.py

@@ -0,0 +1,5 @@
+from .optimizers import *
+
+l1 = ModuleOptimizer.RegularizerType.L1
+l2 = ModuleOptimizer.RegularizerType.L2
+none = ModuleOptimizer.RegularizerType.None_

TensorFrost/sort.py

@@ -0,0 +1,187 @@
+from . import TensorFrost as tf
+
+#in-place bitonic sort
+def bitonic(keys, values = None):
+    tf.region_begin('Bitonic sort')
+    keys = tf.copy(keys)
+    if values is not None:
+        values = tf.copy(values)
+    element_count = keys.shape[0]
+    log2_count = tf.int(tf.ceil(tf.log2(tf.float(element_count))))
+    count_round = 1 << log2_count
+    idx = tf.indices([count_round / 2])[0]
+    with tf.loop(log2_count) as k:
+        with tf.loop(k+1) as j:
+            s = 1 << (k-j)
+            m_inner = s - 1
+            m_outer = ~m_inner
+            m_xor = s + tf.select(j == 0, m_inner, 0)
+
+            id1 = (2 * (idx & m_outer) + (idx & m_inner))
+            id2 = id1 ^ m_xor
+            key1, key2 = keys[id1], keys[id2]
+            with tf.if_cond((key1 >= key2) & (id1 < element_count) & (id2 < element_count)):
+                if values is not None:
+                    val1, val2 = values[id1], values[id2]
+                    values[id1] = val2
+                    values[id2] = val1
+                keys[id1] = key2
+                keys[id2] = key1
+
+    tf.region_end('Bitonic sort')
+    if values is not None:
+        return keys, values
+    else:
+        return keys
+
+#histogram radix sort
+def radix(keys, values = None, bits_per_pass = 6, max_bits = 32):
+    def prefix_sum_grouped(A, axis = -1):
+        axis = len(A.shape) + axis if axis < 0 else axis
+        group_size = 64
+        grouped = tf.split_dim(A, group_size, axis)
+        group_scan = tf.prefix_sum(tf.sum(grouped, axis = axis + 1), axis = axis)
+        ids = grouped.indices
+        gid, eid = ids[axis], ids[axis + 1]
+        ids = [ids[i] for i in range(len(ids)) if i != axis + 1]
+        ids[axis] = gid - 1
+        group_scan = tf.prefix_sum(grouped + tf.select((gid == 0) | (eid != 0), tf.uint(0), group_scan[tuple(ids)]), axis = axis + 1)
+        full_scan = tf.merge_dim(group_scan, target_size = A.shape[axis], axis = axis + 1)
+        return full_scan
+
+    sign_bit = ~tf.uint(0x7FFFFFFF)
+
+    def map_float_to_uint(x):
+        # Convert float to uint representation
+        ux = tf.asuint(x)
+        # Compute mask
+        mask = tf.select((ux >> 31) == 1, ~tf.uint(0), sign_bit)
+        # Apply XOR
+        return ux ^ mask
+
+    def map_uint_to_float(x):
+        # Compute mask
+        mask = tf.select((x >> 31) == 0, ~tf.uint(0), sign_bit)
+        # Apply XOR and convert back to float
+        return tf.asfloat(x ^ mask)
+
+    def map_int_to_uint(x):
+        return tf.asuint(x) ^ sign_bit
+
+    def map_uint_to_int(x):
+        return tf.asint(x ^ sign_bit)
+
+    tf.region_begin('Radix sort')
+
+    has_values = values is not None
+
+    keys = tf.copy(keys)
+    if has_values:
+        values = tf.copy(values)
+
+    original_type = keys.type
+    if(original_type == tf.float32):
+        keys = map_float_to_uint(keys)
+
+    if(original_type == tf.int32):
+        keys = map_int_to_uint(keys)
+
+    iters = (max_bits + bits_per_pass - 1) // bits_per_pass
+    group_size = 128
+    histogram_size = 2 ** bits_per_pass
+
+    def GetBits(A, i):
+        return (A >> (i * bits_per_pass)) & tf.uint(histogram_size - 1)
+
+    keys1 = tf.buffer(keys.shape, keys.type)
+    values1 = None
+
+    if has_values:
+        values1 = tf.buffer(values.shape, values.type)
+
+    with tf.loop(iters // 2) as iter:
+        def SortIteration(keys_in, keys_out, values_in, values_out, iter):
+            tf.region_begin('Radix sort iteration')
+            grouped = tf.split_dim(GetBits(keys_in, iter), group_size)
+
+            # Do a packed histogram, since we sum 128 elements at a time, we can pack 4 values into a single uint32
+            g, e, i = tf.indices([grouped.shape[0], grouped.shape[1], tf.int(histogram_size/4)])
+            this_key = grouped[g, e]
+            packed_is_bit = (tf.uint(this_key == tf.uint(4*i))) + (tf.uint(this_key == tf.uint(4*i+1)) << 8) + (tf.uint(this_key == tf.uint(4*i+2)) << 16) + (tf.uint(this_key == tf.uint(4*i+3)) << 24)
+            packed_is_bit = tf.select((g*group_size + e) < keys_in.shape[0], packed_is_bit, tf.uint(0))
+            group_histogram_packed = tf.sum(packed_is_bit, axis = 1)
+
+            g, i = tf.indices([grouped.shape[0], histogram_size])
+            group_histogram = tf.uint((group_histogram_packed[g, i / 4] >> (8*(i % 4))) & tf.uint(0xFF))
+
+            group_histogram_scan = prefix_sum_grouped(group_histogram, axis = 0)
+            i, = tf.indices([histogram_size])
+            total_bit_histogram = tf.prefix_sum(group_histogram_scan[group_histogram_scan.shape[0] - 1, i])
+
+            with tf.kernel(grouped.shape, group_size=[group_size]) as (g, e):
+                if(tf.current_backend() == tf.cpu): #dont use group barriers on CPU - doesn't work
+                    element = g * group_size + e
+                    with tf.if_cond(element < keys_in.shape[0]):
+                        old_key = keys_in[element]
+                        old_val = values_in[element]
+                        bit = GetBits(old_key, iter)
+                        total_offset = tf.select(g == 0, tf.uint(0), group_histogram_scan[g - 1, bit]) + tf.select(bit == tf.uint(0), tf.uint(0), total_bit_histogram[bit - tf.uint(1)])
+                        with tf.loop(e) as j:
+                            total_offset.val += tf.uint(grouped[g, j] == bit)
+                        keys_out[total_offset] = old_key
+                        values_out[total_offset] = old_val
+                else:
+                    temp = tf.group_buffer(group_size, tf.uint32)
+                    half_count = tf.group_buffer(histogram_size, tf.uint32)
+                    gtid = g.block_thread_index(0)
+
+                    #initialize counters
+                    for i in range((histogram_size + group_size - 1) // group_size):
+                        index = gtid + i * group_size
+                        with tf.if_cond(index < histogram_size):
+                            half_count[index] = 0
+                    tf.group_barrier()
+
+                    element = g * group_size + e
+                    with tf.if_cond(element < keys_in.shape[0]):
+                        old_key = keys_in[element]
+                        bit = GetBits(old_key, iter)
+                        temp[gtid] = bit
+
+                        #count number of bits set in previous sub groups
+                        quarter_index = e / (group_size // 4)
+                        with tf.if_cond(quarter_index < 3):
+                            tf.scatterAdd(half_count[bit], tf.uint(quarter_index < 1) | (tf.uint(quarter_index < 2) << 8) | (tf.uint(quarter_index < 3) << 16))
+
+                        tf.group_barrier()
+
+                        if has_values:
+                            old_val = values_in[element]
+
+                        total_offset = tf.select(g == 0, tf.uint(0), group_histogram_scan[g - 1, tf.int(bit)]) + tf.select(tf.int(bit) == 0, tf.uint(0), total_bit_histogram[tf.int(bit) - 1])
+                        total_offset += tf.select(quarter_index > 0, (half_count[bit] >> (8*(quarter_index-1))) & tf.uint(0xFF), tf.uint(0))
+                        begin_index = quarter_index * (group_size // 4)
+                        with tf.loop(begin_index, e) as j:
+                            total_offset.val += tf.uint(temp[j] == bit)
+                        keys_out[total_offset] = old_key
+
+                        if has_values:
+                            values_out[total_offset] = old_val
+
+            tf.region_end('Radix sort iteration')
+
+        SortIteration(keys, keys1, values, values1, 2 * iter)
+        SortIteration(keys1, keys, values1, values, 2 * iter + 1)
+
+    tf.region_end('Radix sort')
+
+    if(original_type == tf.float32):
+        keys = map_uint_to_float(keys)
+
+    if(original_type == tf.int32):
+        keys = map_uint_to_int(keys)
+
+    if has_values:
+        return keys, values
+    else:
+        return keys