chela 0.0.2

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

import subprocess
import json

import numpy as np
import torch


def merge_dicts(list_of_dicts):
    result = {}
    for k in list_of_dicts[0].keys():
        result[k] = [d[k] for d in list_of_dicts]
    return result


target_to_executable = {}


def compile_rust(target: str) -> str:
    if target not in target_to_executable:
        out = subprocess.check_output(f"$HOME/.cargo/bin/cargo build --release --bin {target} "
                                      "-v --message-format=json", shell=True)
        out = out.split(b"\n")[-3]
        target_to_executable[target] = json.loads(out)["executable"]

    return target_to_executable[target]


def run_rust(executable, *argv) -> float:
    value = subprocess.check_output(f"{executable} {''.join(map(str, argv))}", shell=True)
    return int(value)


def get_class_from_method(method) -> str:
    return method.__qualname__.split(".")[0]


def rand_ndarrays_with_shape(shapes: list[tuple[int, ...]], dtype="float32", slices=None) -> list[np.ndarray]:
    if slices is None:
        return [np.random.rand(*shape).astype(dtype) for shape in shapes]
    return [np.random.rand(*shape).astype(dtype)[slice_] for shape, slice_ in zip(shapes, slices)]


def rand_tensors_with_shape(shapes: list[tuple[int, ...]], dtype=torch.float32, slices=None) -> list:
    if slices is None:
        return [torch.rand(shape).type(dtype) for shape in shapes]
    return [torch.rand(shape).type(dtype)[*slice_] for shape, slice_ in zip(shapes, slices)]


__all__ = ["merge_dicts", "compile_rust", "run_rust", "get_class_from_method",
           "rand_ndarrays_with_shape", "rand_tensors_with_shape"]