redstone-ml 0.0.0

High-performance Machine Learning, Auto-Differentiation and Tensor Algebra crate for Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

from perfprofiler import *

N = 1000000


class TensorBinaryOps(TimingSuite):
    ID: int

    tensors: list
    ndarrays: list[np.ndarray]

    operation = "Addition"

    def __init__(self, shapes, slices=None):
        self.ndarrays = rand_ndarrays_with_shape(shapes, slices=slices, dtype=NUMPY_DTYPE)
        self.tensors = rand_tensors_with_shape(shapes, slices=slices, dtype=TORCH_DTYPE)

    @measure_rust_performance("Chela CPU", target="binary_ops")
    def run(self, executable):
        return run_rust(executable, self.ID, TRIALS, WARMUP)

    @measure_performance("PyTorch CPU")
    def run(self):
        _ = self.tensors[0] + self.tensors[1]

    @measure_performance("NumPy")
    def run(self):
        _ = self.ndarrays[0] + self.ndarrays[1]


class TensorBinaryOps0(TensorBinaryOps):
    ID = 0
    name = "[1], [1]"

    def __init__(self):
        super().__init__([(N, ), (N, )])


class TensorBinaryOps1(TensorBinaryOps):
    ID = 1
    name = "[1], [0]"

    def __init__(self):
        super().__init__([(N,), (1,)])


class TensorBinaryOps2(TensorBinaryOps):
    ID = 2
    name = "[n], [1]"

    def __init__(self):
        super().__init__([(N, 3), (N, )], slices=((Ellipsis, 0), ()))


class TensorBinaryOps3(TensorBinaryOps):
    ID = 2
    name = "[n], [0]"

    def __init__(self):
        super().__init__([(N, 3), (1, )], slices=((Ellipsis, 0), ()))


class TensorBinaryOps4(TensorBinaryOps):
    ID = 4
    name = "Non-Unif, [1]"

    def __init__(self):
        super().__init__([(N, 3), (N, 2)], slices=((Ellipsis, slice(0, 2)), ()))


class TensorBinaryOps5(TensorBinaryOps):
    ID = 5
    name = "Non-Unif, [0]"

    def __init__(self):
        super().__init__([(N, 3), (1,)], slices=((Ellipsis, slice(0, 2)), ()))


class TensorBinaryOps6(TensorBinaryOps):
    ID = 6
    name = "Non-Unif, [n]"

    def __init__(self):
        super().__init__([(N, 3), (N, 2, 2)], slices=((Ellipsis, slice(0, 2)), (Ellipsis, 0)))


class TensorBinaryOps7(TensorBinaryOps):
    ID = 7
    name = "Non-Unif, Non-Unif"

    def __init__(self):
        super().__init__([(N, 3), (N, 3)], slices=((Ellipsis, slice(0, 2)), (Ellipsis, slice(0, 2))))


class TensorBinaryOps8(TensorBinaryOps):
    ID = 8
    name = "[n], [n]"

    def __init__(self):
        super().__init__([(N, 3), (N, 3)], slices=((Ellipsis, 0), (Ellipsis, 0)))


if __name__ == "__main__":
    results = profile_all([
        TensorBinaryOps0,
        TensorBinaryOps1,
        TensorBinaryOps2,
        TensorBinaryOps8,

        TensorBinaryOps3,
        TensorBinaryOps4,
        TensorBinaryOps5,
        TensorBinaryOps6,
        TensorBinaryOps7
    ])
    plot_barplot(results, f"{TensorBinaryOps.operation} Benchmark ({N=})", normalize="NumPy")