modelstudio 0.1.1__tar.gz → 0.3.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {modelstudio-0.1.1 → modelstudio-0.3.0}/MANIFEST.in +2 -0
- modelstudio-0.3.0/PKG-INFO +223 -0
- modelstudio-0.3.0/README.md +192 -0
- modelstudio-0.3.0/benchmarks/bench_attention.py +39 -0
- modelstudio-0.3.0/benchmarks/bench_conv.py +40 -0
- modelstudio-0.3.0/benchmarks/bench_dataloader.py +46 -0
- modelstudio-0.3.0/benchmarks/bench_dropout.py +37 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/CMakeLists.txt +2 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/backends/cpu/cpu_backend.cpp +8 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/backends/cpu/cpu_backend.hpp +4 -0
- modelstudio-0.3.0/csrc/backends/cpu/kernels/mul.cpp +43 -0
- modelstudio-0.3.0/csrc/backends/cpu/kernels/relu.cpp +40 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/backends/cuda/cuda_backend.cu +10 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/backends/cuda/cuda_backend.hpp +2 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/backends/oneapi/oneapi_backend.cpp +10 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/backends/oneapi/oneapi_backend.hpp +2 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/backends/rocm/rocm_backend.cpp +10 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/backends/rocm/rocm_backend.hpp +2 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/dispatcher/backend.hpp +2 -0
- modelstudio-0.3.0/docs/autograd.md +55 -0
- modelstudio-0.3.0/docs/backend-architecture.md +72 -0
- modelstudio-0.3.0/docs/data.md +33 -0
- modelstudio-0.3.0/docs/modules.md +24 -0
- modelstudio-0.3.0/docs/native-backend-roadmap.md +26 -0
- modelstudio-0.3.0/docs/nn.md +68 -0
- modelstudio-0.3.0/docs/randomness.md +23 -0
- modelstudio-0.3.0/docs/releasing.md +86 -0
- modelstudio-0.3.0/docs/serialization.md +25 -0
- modelstudio-0.3.0/docs/tensor-api.md +51 -0
- modelstudio-0.3.0/docs/training.md +29 -0
- modelstudio-0.3.0/examples/checkpoint_training.py +39 -0
- modelstudio-0.3.0/examples/dropout_batchnorm.py +35 -0
- modelstudio-0.3.0/examples/save_load.py +21 -0
- modelstudio-0.3.0/examples/tiny_transformer.py +35 -0
- modelstudio-0.3.0/examples/train_classifier.py +28 -0
- modelstudio-0.3.0/examples/train_cnn_toy.py +40 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/pyproject.toml +1 -1
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/__init__.py +27 -4
- modelstudio-0.3.0/python/modelstudio/_version.py +1 -0
- modelstudio-0.3.0/python/modelstudio/data/__init__.py +4 -0
- modelstudio-0.3.0/python/modelstudio/data/dataloader.py +68 -0
- modelstudio-0.3.0/python/modelstudio/data/dataset.py +30 -0
- modelstudio-0.3.0/python/modelstudio/nn/__init__.py +36 -0
- modelstudio-0.3.0/python/modelstudio/nn/activations.py +29 -0
- modelstudio-0.3.0/python/modelstudio/nn/convolution.py +213 -0
- modelstudio-0.3.0/python/modelstudio/nn/embedding.py +19 -0
- modelstudio-0.3.0/python/modelstudio/nn/init.py +57 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/nn/linear.py +8 -5
- modelstudio-0.3.0/python/modelstudio/nn/losses.py +83 -0
- modelstudio-0.3.0/python/modelstudio/nn/module.py +239 -0
- modelstudio-0.3.0/python/modelstudio/nn/normalization.py +93 -0
- modelstudio-0.3.0/python/modelstudio/nn/pooling.py +135 -0
- modelstudio-0.3.0/python/modelstudio/nn/transformer.py +41 -0
- modelstudio-0.3.0/python/modelstudio/nn/utils.py +33 -0
- modelstudio-0.3.0/python/modelstudio/ops/__init__.py +58 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/ops/linalg.py +19 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/ops/math.py +168 -1
- modelstudio-0.3.0/python/modelstudio/ops/movement.py +214 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/ops/reductions.py +6 -0
- modelstudio-0.3.0/python/modelstudio/optim/adamw.py +120 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/optim/optimizer.py +14 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/optim/sgd.py +18 -0
- modelstudio-0.3.0/python/modelstudio/random.py +20 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/runtime/backend.py +2 -1
- modelstudio-0.3.0/python/modelstudio/serialization.py +63 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/tensor.py +64 -0
- modelstudio-0.3.0/python/modelstudio.egg-info/PKG-INFO +223 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio.egg-info/SOURCES.txt +60 -0
- modelstudio-0.3.0/scripts/smoke_test.py +51 -0
- modelstudio-0.3.0/tests/test_attention.py +27 -0
- modelstudio-0.3.0/tests/test_batchnorm.py +50 -0
- modelstudio-0.3.0/tests/test_buffers.py +39 -0
- modelstudio-0.3.0/tests/test_concat_stack.py +51 -0
- modelstudio-0.3.0/tests/test_conv.py +57 -0
- modelstudio-0.3.0/tests/test_data.py +42 -0
- modelstudio-0.3.0/tests/test_dataloader_seed.py +44 -0
- modelstudio-0.3.0/tests/test_dropout.py +60 -0
- modelstudio-0.3.0/tests/test_embedding.py +22 -0
- modelstudio-0.3.0/tests/test_grad_clip.py +24 -0
- modelstudio-0.3.0/tests/test_indexing.py +34 -0
- modelstudio-0.3.0/tests/test_init.py +43 -0
- modelstudio-0.3.0/tests/test_loss_reductions.py +34 -0
- modelstudio-0.3.0/tests/test_losses.py +36 -0
- modelstudio-0.3.0/tests/test_module_ergonomics.py +85 -0
- modelstudio-0.3.0/tests/test_norms.py +31 -0
- modelstudio-0.3.0/tests/test_optimizer_state.py +89 -0
- modelstudio-0.3.0/tests/test_pooling.py +31 -0
- modelstudio-0.3.0/tests/test_random.py +23 -0
- modelstudio-0.3.0/tests/test_reductions_axis.py +52 -0
- modelstudio-0.3.0/tests/test_serialization.py +38 -0
- modelstudio-0.3.0/tests/test_shape_ops.py +33 -0
- modelstudio-0.3.0/tests/test_state_dict.py +71 -0
- modelstudio-0.3.0/tests/test_transformer.py +26 -0
- modelstudio-0.3.0/tests/test_unary_ops.py +47 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/tests/test_version.py +1 -1
- modelstudio-0.1.1/PKG-INFO +0 -235
- modelstudio-0.1.1/README.md +0 -204
- modelstudio-0.1.1/python/modelstudio/_version.py +0 -1
- modelstudio-0.1.1/python/modelstudio/nn/__init__.py +0 -8
- modelstudio-0.1.1/python/modelstudio/nn/activations.py +0 -14
- modelstudio-0.1.1/python/modelstudio/nn/losses.py +0 -14
- modelstudio-0.1.1/python/modelstudio/nn/module.py +0 -76
- modelstudio-0.1.1/python/modelstudio/nn/normalization.py +0 -27
- modelstudio-0.1.1/python/modelstudio/ops/__init__.py +0 -24
- modelstudio-0.1.1/python/modelstudio/ops/movement.py +0 -58
- modelstudio-0.1.1/python/modelstudio/optim/adamw.py +0 -55
- modelstudio-0.1.1/python/modelstudio.egg-info/PKG-INFO +0 -235
- {modelstudio-0.1.1 → modelstudio-0.3.0}/CMakeLists.txt +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/LICENSE +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/benchmarks/bench_matmul.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/benchmarks/bench_mlp.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/backends/cpu/kernels/add.cpp +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/backends/cpu/kernels/matmul.cpp +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/backends/cuda/README.md +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/backends/cuda/cuda_memory.hpp +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/backends/oneapi/README.md +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/backends/oneapi/sycl_memory.hpp +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/backends/rocm/README.md +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/backends/rocm/hip_memory.hpp +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/bindings/python_bindings.cpp +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/core/device.hpp +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/core/dtype.hpp +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/core/error.hpp +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/core/shape.hpp +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/core/storage.hpp +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/core/tensor.hpp +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/dispatcher/dispatcher.hpp +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/csrc/dispatcher/operator_registry.hpp +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/examples/train_mlp.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/autograd/__init__.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/autograd/engine.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/autograd/function.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/autograd/grad_mode.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/compile/__init__.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/compile/graph_capture.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/compile/ir.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/compile/passes.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/device.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/dtypes.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/errors.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/nn/parameter.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/ops/creation.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/optim/__init__.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/py.typed +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/runtime/__init__.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/runtime/dispatcher.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/storage.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/testing/__init__.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio/testing/gradcheck.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio.egg-info/dependency_links.txt +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio.egg-info/requires.txt +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/python/modelstudio.egg-info/top_level.txt +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/setup.cfg +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/tests/test_autograd.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/tests/test_dispatcher.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/tests/test_gradcheck.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/tests/test_nn.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/tests/test_ops.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/tests/test_optim.py +0 -0
- {modelstudio-0.1.1 → modelstudio-0.3.0}/tests/test_tensor.py +0 -0
|
@@ -0,0 +1,223 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: modelstudio
|
|
3
|
+
Version: 0.3.0
|
|
4
|
+
Summary: An early-stage AI tensor framework with CPU tensors, autograd, and backend extension scaffolding.
|
|
5
|
+
Author: ModelStudio Contributors
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/imattas/modelstudio
|
|
8
|
+
Project-URL: Repository, https://github.com/imattas/modelstudio
|
|
9
|
+
Project-URL: Issues, https://github.com/imattas/modelstudio/issues
|
|
10
|
+
Keywords: ai,autograd,deep-learning,neural-networks,tensor
|
|
11
|
+
Classifier: Development Status :: 3 - Alpha
|
|
12
|
+
Classifier: Intended Audience :: Developers
|
|
13
|
+
Classifier: Intended Audience :: Science/Research
|
|
14
|
+
Classifier: Programming Language :: Python :: 3
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
19
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
20
|
+
Classifier: Typing :: Typed
|
|
21
|
+
Requires-Python: >=3.10
|
|
22
|
+
Description-Content-Type: text/markdown
|
|
23
|
+
License-File: LICENSE
|
|
24
|
+
Requires-Dist: numpy>=1.26
|
|
25
|
+
Provides-Extra: dev
|
|
26
|
+
Requires-Dist: pytest>=8; extra == "dev"
|
|
27
|
+
Requires-Dist: ruff>=0.6; extra == "dev"
|
|
28
|
+
Requires-Dist: build>=1.2; extra == "dev"
|
|
29
|
+
Requires-Dist: twine>=5; extra == "dev"
|
|
30
|
+
Dynamic: license-file
|
|
31
|
+
|
|
32
|
+
# ModelStudio
|
|
33
|
+
|
|
34
|
+
ModelStudio is an early-stage AI tensor framework. Version `0.3.0` provides a
|
|
35
|
+
CPU tensor/autograd MVP with neural-network modules, optimizers, serialization,
|
|
36
|
+
basic data loading, and small LLM-oriented building blocks.
|
|
37
|
+
|
|
38
|
+
It is not a PyTorch or TensorFlow replacement. CPU is the only working backend.
|
|
39
|
+
CUDA, ROCm, and oneAPI remain explicit scaffolds until real kernels are built
|
|
40
|
+
and tested.
|
|
41
|
+
|
|
42
|
+
## Installation
|
|
43
|
+
|
|
44
|
+
From PyPI:
|
|
45
|
+
|
|
46
|
+
```bash
|
|
47
|
+
python -m pip install modelstudio
|
|
48
|
+
```
|
|
49
|
+
|
|
50
|
+
For development:
|
|
51
|
+
|
|
52
|
+
```bash
|
|
53
|
+
python -m pip install -e ".[dev]"
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
## Feature Table
|
|
57
|
+
|
|
58
|
+
| Area | Status |
|
|
59
|
+
| --- | --- |
|
|
60
|
+
| CPU tensors | Working MVP |
|
|
61
|
+
| Autograd | Reverse-mode for core CPU ops |
|
|
62
|
+
| Reductions | `sum`, `mean`, `max` with axis and keepdims; `max` is value-only |
|
|
63
|
+
| Activations | ReLU, GELU, exp, log, tanh, sigmoid, SiLU, softmax, log-softmax |
|
|
64
|
+
| Losses | MSE and cross entropy with `none`, `mean`, and `sum` reductions |
|
|
65
|
+
| Modules | Parameters, buffers, child traversal, state dicts, save/load |
|
|
66
|
+
| Layers | Linear, Embedding, LayerNorm, RMSNorm, BatchNorm1d, Dropout, Conv1d, Conv2d, pooling, TransformerBlock |
|
|
67
|
+
| Optimizers | SGD and AdamW with state serialization |
|
|
68
|
+
| Data | Dataset, TensorDataset, DataLoader with deterministic seeded shuffle |
|
|
69
|
+
| Randomness | `manual_seed`, RNG-backed `randn`, dropout, and init helpers |
|
|
70
|
+
| Compiler | Placeholder IR and passes |
|
|
71
|
+
|
|
72
|
+
## Backend Status
|
|
73
|
+
|
|
74
|
+
| Backend | Status |
|
|
75
|
+
| --- | --- |
|
|
76
|
+
| CPU | working MVP |
|
|
77
|
+
| CUDA | scaffold only |
|
|
78
|
+
| ROCm | scaffold only |
|
|
79
|
+
| oneAPI | scaffold only |
|
|
80
|
+
|
|
81
|
+
Unsupported accelerator devices fail with `ModelStudioBackendUnavailable`.
|
|
82
|
+
|
|
83
|
+
## Tensor Example
|
|
84
|
+
|
|
85
|
+
```python
|
|
86
|
+
import modelstudio as ms
|
|
87
|
+
|
|
88
|
+
x = ms.randn((32, 784), requires_grad=True)
|
|
89
|
+
w = ms.randn((784, 10), requires_grad=True)
|
|
90
|
+
loss = (x @ w).mean()
|
|
91
|
+
loss.backward()
|
|
92
|
+
print(w.grad)
|
|
93
|
+
```
|
|
94
|
+
|
|
95
|
+
## MLP Example
|
|
96
|
+
|
|
97
|
+
```python
|
|
98
|
+
import modelstudio as ms
|
|
99
|
+
from modelstudio import nn
|
|
100
|
+
|
|
101
|
+
|
|
102
|
+
class MLP(nn.Module):
|
|
103
|
+
def __init__(self):
|
|
104
|
+
super().__init__()
|
|
105
|
+
self.fc1 = nn.Linear(784, 256)
|
|
106
|
+
self.fc2 = nn.Linear(256, 10)
|
|
107
|
+
|
|
108
|
+
def forward(self, x):
|
|
109
|
+
return self.fc2(ms.gelu(self.fc1(x)))
|
|
110
|
+
|
|
111
|
+
|
|
112
|
+
model = MLP()
|
|
113
|
+
optimizer = ms.optim.AdamW(model.parameters(), lr=3e-4)
|
|
114
|
+
x = ms.randn((16, 784))
|
|
115
|
+
target = ms.randn((16, 10))
|
|
116
|
+
loss = ms.mse_loss(model(x), target)
|
|
117
|
+
optimizer.zero_grad()
|
|
118
|
+
loss.backward()
|
|
119
|
+
optimizer.step()
|
|
120
|
+
```
|
|
121
|
+
|
|
122
|
+
## State Dict and Save/Load
|
|
123
|
+
|
|
124
|
+
```python
|
|
125
|
+
model = nn.Linear(4, 2)
|
|
126
|
+
ms.save(model.state_dict(), "model.ms")
|
|
127
|
+
state = ms.load("model.ms")
|
|
128
|
+
model.load_state_dict(state)
|
|
129
|
+
```
|
|
130
|
+
|
|
131
|
+
## DataLoader
|
|
132
|
+
|
|
133
|
+
```python
|
|
134
|
+
from modelstudio import data
|
|
135
|
+
|
|
136
|
+
dataset = data.TensorDataset(ms.randn((8, 4)), ms.arange(8))
|
|
137
|
+
loader = data.DataLoader(dataset, batch_size=2, shuffle=False)
|
|
138
|
+
for xb, yb in loader:
|
|
139
|
+
print(xb.shape, yb.shape)
|
|
140
|
+
```
|
|
141
|
+
|
|
142
|
+
## Embedding
|
|
143
|
+
|
|
144
|
+
```python
|
|
145
|
+
emb = nn.Embedding(num_embeddings=100, embedding_dim=32)
|
|
146
|
+
tokens = ms.tensor([[1, 2, 3]], dtype=ms.int64)
|
|
147
|
+
print(emb(tokens).shape)
|
|
148
|
+
```
|
|
149
|
+
|
|
150
|
+
## Cross Entropy
|
|
151
|
+
|
|
152
|
+
```python
|
|
153
|
+
logits = ms.randn((4, 10), requires_grad=True)
|
|
154
|
+
targets = ms.tensor([1, 2, 3, 4], dtype=ms.int64)
|
|
155
|
+
loss = ms.cross_entropy(logits, targets)
|
|
156
|
+
loss.backward()
|
|
157
|
+
```
|
|
158
|
+
|
|
159
|
+
## TransformerBlock
|
|
160
|
+
|
|
161
|
+
```python
|
|
162
|
+
block = nn.TransformerBlock(embed_dim=16, num_heads=4)
|
|
163
|
+
x = ms.randn((2, 8, 16), requires_grad=True)
|
|
164
|
+
y = block(x)
|
|
165
|
+
print(y.shape)
|
|
166
|
+
```
|
|
167
|
+
|
|
168
|
+
## 0.3.0 Training Utilities
|
|
169
|
+
|
|
170
|
+
```python
|
|
171
|
+
ms.manual_seed(123)
|
|
172
|
+
model = nn.Linear(4, 2)
|
|
173
|
+
optimizer = ms.optim.AdamW(model.parameters(), lr=1e-3)
|
|
174
|
+
state = {"model": model.state_dict(), "optimizer": optimizer.state_dict()}
|
|
175
|
+
ms.save(state, "checkpoint.ms")
|
|
176
|
+
```
|
|
177
|
+
|
|
178
|
+
New CPU-only helpers include `ms.concat`, `ms.stack`, `Tensor.flatten`,
|
|
179
|
+
`Tensor.squeeze`, `Tensor.unsqueeze`, `nn.init`, `nn.Dropout`,
|
|
180
|
+
`nn.BatchNorm1d`, `nn.Conv1d`, `nn.Conv2d`, `nn.AvgPool2d`, `nn.MaxPool2d`,
|
|
181
|
+
and `nn.utils` gradient clipping.
|
|
182
|
+
|
|
183
|
+
## Commands
|
|
184
|
+
|
|
185
|
+
```bash
|
|
186
|
+
python -m pytest
|
|
187
|
+
python scripts/smoke_test.py
|
|
188
|
+
python examples/train_mlp.py
|
|
189
|
+
python examples/train_classifier.py
|
|
190
|
+
python examples/tiny_transformer.py
|
|
191
|
+
python examples/save_load.py
|
|
192
|
+
python examples/train_cnn_toy.py
|
|
193
|
+
python examples/dropout_batchnorm.py
|
|
194
|
+
python examples/checkpoint_training.py
|
|
195
|
+
python benchmarks/bench_matmul.py
|
|
196
|
+
python benchmarks/bench_mlp.py
|
|
197
|
+
python benchmarks/bench_attention.py
|
|
198
|
+
python benchmarks/bench_dataloader.py
|
|
199
|
+
python benchmarks/bench_conv.py
|
|
200
|
+
python benchmarks/bench_dropout.py
|
|
201
|
+
```
|
|
202
|
+
|
|
203
|
+
## Documentation
|
|
204
|
+
|
|
205
|
+
- [Tensor API](docs/tensor-api.md)
|
|
206
|
+
- [Neural network API](docs/nn.md)
|
|
207
|
+
- [Data utilities](docs/data.md)
|
|
208
|
+
- [Training](docs/training.md)
|
|
209
|
+
- [Modules](docs/modules.md)
|
|
210
|
+
- [Serialization](docs/serialization.md)
|
|
211
|
+
- [Randomness](docs/randomness.md)
|
|
212
|
+
- [Native backend roadmap](docs/native-backend-roadmap.md)
|
|
213
|
+
- [Backend architecture](docs/backend-architecture.md)
|
|
214
|
+
- [Autograd design](docs/autograd.md)
|
|
215
|
+
- [Releasing](docs/releasing.md)
|
|
216
|
+
- [Contributing](CONTRIBUTING.md)
|
|
217
|
+
|
|
218
|
+
## Roadmap
|
|
219
|
+
|
|
220
|
+
- Expand tensor and autograd coverage.
|
|
221
|
+
- Wire native CPU kernels into Python bindings.
|
|
222
|
+
- Add tested CUDA, ROCm, and oneAPI packages when hardware-backed CI exists.
|
|
223
|
+
- Improve compiler graph capture and lowering.
|
|
@@ -0,0 +1,192 @@
|
|
|
1
|
+
# ModelStudio
|
|
2
|
+
|
|
3
|
+
ModelStudio is an early-stage AI tensor framework. Version `0.3.0` provides a
|
|
4
|
+
CPU tensor/autograd MVP with neural-network modules, optimizers, serialization,
|
|
5
|
+
basic data loading, and small LLM-oriented building blocks.
|
|
6
|
+
|
|
7
|
+
It is not a PyTorch or TensorFlow replacement. CPU is the only working backend.
|
|
8
|
+
CUDA, ROCm, and oneAPI remain explicit scaffolds until real kernels are built
|
|
9
|
+
and tested.
|
|
10
|
+
|
|
11
|
+
## Installation
|
|
12
|
+
|
|
13
|
+
From PyPI:
|
|
14
|
+
|
|
15
|
+
```bash
|
|
16
|
+
python -m pip install modelstudio
|
|
17
|
+
```
|
|
18
|
+
|
|
19
|
+
For development:
|
|
20
|
+
|
|
21
|
+
```bash
|
|
22
|
+
python -m pip install -e ".[dev]"
|
|
23
|
+
```
|
|
24
|
+
|
|
25
|
+
## Feature Table
|
|
26
|
+
|
|
27
|
+
| Area | Status |
|
|
28
|
+
| --- | --- |
|
|
29
|
+
| CPU tensors | Working MVP |
|
|
30
|
+
| Autograd | Reverse-mode for core CPU ops |
|
|
31
|
+
| Reductions | `sum`, `mean`, `max` with axis and keepdims; `max` is value-only |
|
|
32
|
+
| Activations | ReLU, GELU, exp, log, tanh, sigmoid, SiLU, softmax, log-softmax |
|
|
33
|
+
| Losses | MSE and cross entropy with `none`, `mean`, and `sum` reductions |
|
|
34
|
+
| Modules | Parameters, buffers, child traversal, state dicts, save/load |
|
|
35
|
+
| Layers | Linear, Embedding, LayerNorm, RMSNorm, BatchNorm1d, Dropout, Conv1d, Conv2d, pooling, TransformerBlock |
|
|
36
|
+
| Optimizers | SGD and AdamW with state serialization |
|
|
37
|
+
| Data | Dataset, TensorDataset, DataLoader with deterministic seeded shuffle |
|
|
38
|
+
| Randomness | `manual_seed`, RNG-backed `randn`, dropout, and init helpers |
|
|
39
|
+
| Compiler | Placeholder IR and passes |
|
|
40
|
+
|
|
41
|
+
## Backend Status
|
|
42
|
+
|
|
43
|
+
| Backend | Status |
|
|
44
|
+
| --- | --- |
|
|
45
|
+
| CPU | working MVP |
|
|
46
|
+
| CUDA | scaffold only |
|
|
47
|
+
| ROCm | scaffold only |
|
|
48
|
+
| oneAPI | scaffold only |
|
|
49
|
+
|
|
50
|
+
Unsupported accelerator devices fail with `ModelStudioBackendUnavailable`.
|
|
51
|
+
|
|
52
|
+
## Tensor Example
|
|
53
|
+
|
|
54
|
+
```python
|
|
55
|
+
import modelstudio as ms
|
|
56
|
+
|
|
57
|
+
x = ms.randn((32, 784), requires_grad=True)
|
|
58
|
+
w = ms.randn((784, 10), requires_grad=True)
|
|
59
|
+
loss = (x @ w).mean()
|
|
60
|
+
loss.backward()
|
|
61
|
+
print(w.grad)
|
|
62
|
+
```
|
|
63
|
+
|
|
64
|
+
## MLP Example
|
|
65
|
+
|
|
66
|
+
```python
|
|
67
|
+
import modelstudio as ms
|
|
68
|
+
from modelstudio import nn
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
class MLP(nn.Module):
|
|
72
|
+
def __init__(self):
|
|
73
|
+
super().__init__()
|
|
74
|
+
self.fc1 = nn.Linear(784, 256)
|
|
75
|
+
self.fc2 = nn.Linear(256, 10)
|
|
76
|
+
|
|
77
|
+
def forward(self, x):
|
|
78
|
+
return self.fc2(ms.gelu(self.fc1(x)))
|
|
79
|
+
|
|
80
|
+
|
|
81
|
+
model = MLP()
|
|
82
|
+
optimizer = ms.optim.AdamW(model.parameters(), lr=3e-4)
|
|
83
|
+
x = ms.randn((16, 784))
|
|
84
|
+
target = ms.randn((16, 10))
|
|
85
|
+
loss = ms.mse_loss(model(x), target)
|
|
86
|
+
optimizer.zero_grad()
|
|
87
|
+
loss.backward()
|
|
88
|
+
optimizer.step()
|
|
89
|
+
```
|
|
90
|
+
|
|
91
|
+
## State Dict and Save/Load
|
|
92
|
+
|
|
93
|
+
```python
|
|
94
|
+
model = nn.Linear(4, 2)
|
|
95
|
+
ms.save(model.state_dict(), "model.ms")
|
|
96
|
+
state = ms.load("model.ms")
|
|
97
|
+
model.load_state_dict(state)
|
|
98
|
+
```
|
|
99
|
+
|
|
100
|
+
## DataLoader
|
|
101
|
+
|
|
102
|
+
```python
|
|
103
|
+
from modelstudio import data
|
|
104
|
+
|
|
105
|
+
dataset = data.TensorDataset(ms.randn((8, 4)), ms.arange(8))
|
|
106
|
+
loader = data.DataLoader(dataset, batch_size=2, shuffle=False)
|
|
107
|
+
for xb, yb in loader:
|
|
108
|
+
print(xb.shape, yb.shape)
|
|
109
|
+
```
|
|
110
|
+
|
|
111
|
+
## Embedding
|
|
112
|
+
|
|
113
|
+
```python
|
|
114
|
+
emb = nn.Embedding(num_embeddings=100, embedding_dim=32)
|
|
115
|
+
tokens = ms.tensor([[1, 2, 3]], dtype=ms.int64)
|
|
116
|
+
print(emb(tokens).shape)
|
|
117
|
+
```
|
|
118
|
+
|
|
119
|
+
## Cross Entropy
|
|
120
|
+
|
|
121
|
+
```python
|
|
122
|
+
logits = ms.randn((4, 10), requires_grad=True)
|
|
123
|
+
targets = ms.tensor([1, 2, 3, 4], dtype=ms.int64)
|
|
124
|
+
loss = ms.cross_entropy(logits, targets)
|
|
125
|
+
loss.backward()
|
|
126
|
+
```
|
|
127
|
+
|
|
128
|
+
## TransformerBlock
|
|
129
|
+
|
|
130
|
+
```python
|
|
131
|
+
block = nn.TransformerBlock(embed_dim=16, num_heads=4)
|
|
132
|
+
x = ms.randn((2, 8, 16), requires_grad=True)
|
|
133
|
+
y = block(x)
|
|
134
|
+
print(y.shape)
|
|
135
|
+
```
|
|
136
|
+
|
|
137
|
+
## 0.3.0 Training Utilities
|
|
138
|
+
|
|
139
|
+
```python
|
|
140
|
+
ms.manual_seed(123)
|
|
141
|
+
model = nn.Linear(4, 2)
|
|
142
|
+
optimizer = ms.optim.AdamW(model.parameters(), lr=1e-3)
|
|
143
|
+
state = {"model": model.state_dict(), "optimizer": optimizer.state_dict()}
|
|
144
|
+
ms.save(state, "checkpoint.ms")
|
|
145
|
+
```
|
|
146
|
+
|
|
147
|
+
New CPU-only helpers include `ms.concat`, `ms.stack`, `Tensor.flatten`,
|
|
148
|
+
`Tensor.squeeze`, `Tensor.unsqueeze`, `nn.init`, `nn.Dropout`,
|
|
149
|
+
`nn.BatchNorm1d`, `nn.Conv1d`, `nn.Conv2d`, `nn.AvgPool2d`, `nn.MaxPool2d`,
|
|
150
|
+
and `nn.utils` gradient clipping.
|
|
151
|
+
|
|
152
|
+
## Commands
|
|
153
|
+
|
|
154
|
+
```bash
|
|
155
|
+
python -m pytest
|
|
156
|
+
python scripts/smoke_test.py
|
|
157
|
+
python examples/train_mlp.py
|
|
158
|
+
python examples/train_classifier.py
|
|
159
|
+
python examples/tiny_transformer.py
|
|
160
|
+
python examples/save_load.py
|
|
161
|
+
python examples/train_cnn_toy.py
|
|
162
|
+
python examples/dropout_batchnorm.py
|
|
163
|
+
python examples/checkpoint_training.py
|
|
164
|
+
python benchmarks/bench_matmul.py
|
|
165
|
+
python benchmarks/bench_mlp.py
|
|
166
|
+
python benchmarks/bench_attention.py
|
|
167
|
+
python benchmarks/bench_dataloader.py
|
|
168
|
+
python benchmarks/bench_conv.py
|
|
169
|
+
python benchmarks/bench_dropout.py
|
|
170
|
+
```
|
|
171
|
+
|
|
172
|
+
## Documentation
|
|
173
|
+
|
|
174
|
+
- [Tensor API](docs/tensor-api.md)
|
|
175
|
+
- [Neural network API](docs/nn.md)
|
|
176
|
+
- [Data utilities](docs/data.md)
|
|
177
|
+
- [Training](docs/training.md)
|
|
178
|
+
- [Modules](docs/modules.md)
|
|
179
|
+
- [Serialization](docs/serialization.md)
|
|
180
|
+
- [Randomness](docs/randomness.md)
|
|
181
|
+
- [Native backend roadmap](docs/native-backend-roadmap.md)
|
|
182
|
+
- [Backend architecture](docs/backend-architecture.md)
|
|
183
|
+
- [Autograd design](docs/autograd.md)
|
|
184
|
+
- [Releasing](docs/releasing.md)
|
|
185
|
+
- [Contributing](CONTRIBUTING.md)
|
|
186
|
+
|
|
187
|
+
## Roadmap
|
|
188
|
+
|
|
189
|
+
- Expand tensor and autograd coverage.
|
|
190
|
+
- Wire native CPU kernels into Python bindings.
|
|
191
|
+
- Add tested CUDA, ROCm, and oneAPI packages when hardware-backed CI exists.
|
|
192
|
+
- Improve compiler graph capture and lowering.
|
|
@@ -0,0 +1,39 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import platform
|
|
4
|
+
import time
|
|
5
|
+
|
|
6
|
+
import modelstudio as ms
|
|
7
|
+
import numpy as np
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
def timeit(fn, iterations: int = 20, warmup: int = 5) -> float:
|
|
11
|
+
for _ in range(warmup):
|
|
12
|
+
fn()
|
|
13
|
+
start = time.perf_counter()
|
|
14
|
+
for _ in range(iterations):
|
|
15
|
+
fn()
|
|
16
|
+
return (time.perf_counter() - start) / iterations
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
def main() -> None:
|
|
20
|
+
shape = (2, 4, 16, 8)
|
|
21
|
+
iterations = 20
|
|
22
|
+
warmup = 5
|
|
23
|
+
q = ms.randn(shape)
|
|
24
|
+
k = ms.randn(shape)
|
|
25
|
+
v = ms.randn(shape)
|
|
26
|
+
|
|
27
|
+
print(f"Python: {platform.python_version()}")
|
|
28
|
+
print(f"NumPy: {np.__version__}")
|
|
29
|
+
print(f"ModelStudio: {ms.__version__}")
|
|
30
|
+
print(f"Shape: {shape}")
|
|
31
|
+
print(f"Warmup: {warmup}")
|
|
32
|
+
print(f"Iterations: {iterations}")
|
|
33
|
+
print("Backend: CPU MVP")
|
|
34
|
+
elapsed = timeit(lambda: ms.scaled_dot_product_attention(q, k, v, causal=True), iterations, warmup)
|
|
35
|
+
print(f"Attention avg: {elapsed * 1_000:.3f} ms")
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
if __name__ == "__main__":
|
|
39
|
+
main()
|
|
@@ -0,0 +1,40 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import platform
|
|
4
|
+
import time
|
|
5
|
+
|
|
6
|
+
import modelstudio as ms
|
|
7
|
+
import numpy as np
|
|
8
|
+
from modelstudio import nn
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
def timeit(fn, iterations: int, warmup: int) -> float:
|
|
12
|
+
for _ in range(warmup):
|
|
13
|
+
fn()
|
|
14
|
+
start = time.perf_counter()
|
|
15
|
+
for _ in range(iterations):
|
|
16
|
+
fn()
|
|
17
|
+
return (time.perf_counter() - start) / iterations
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def main() -> None:
|
|
21
|
+
warmup = 3
|
|
22
|
+
iterations = 10
|
|
23
|
+
shape = (8, 3, 16, 16)
|
|
24
|
+
ms.manual_seed(1)
|
|
25
|
+
x = ms.randn(shape)
|
|
26
|
+
conv = nn.Conv2d(3, 8, kernel_size=3, padding=1)
|
|
27
|
+
|
|
28
|
+
print(f"Python: {platform.python_version()}")
|
|
29
|
+
print(f"NumPy: {np.__version__}")
|
|
30
|
+
print(f"ModelStudio: {ms.__version__}")
|
|
31
|
+
print(f"Operation: Conv2d input={shape} out_channels=8 kernel=3 padding=1")
|
|
32
|
+
print(f"Warmup: {warmup}")
|
|
33
|
+
print(f"Iterations: {iterations}")
|
|
34
|
+
print("Backend: CPU only")
|
|
35
|
+
print(f"Conv2d avg: {timeit(lambda: conv(x), iterations, warmup) * 1_000:.3f} ms")
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
if __name__ == "__main__":
|
|
39
|
+
main()
|
|
40
|
+
|
|
@@ -0,0 +1,46 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import platform
|
|
4
|
+
import time
|
|
5
|
+
|
|
6
|
+
import modelstudio as ms
|
|
7
|
+
import numpy as np
|
|
8
|
+
from modelstudio import data
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
def main() -> None:
|
|
12
|
+
samples = 1024
|
|
13
|
+
features = 32
|
|
14
|
+
batch_size = 64
|
|
15
|
+
iterations = 20
|
|
16
|
+
warmup = 5
|
|
17
|
+
dataset = data.TensorDataset(ms.randn((samples, features)), ms.arange(samples))
|
|
18
|
+
loader = data.DataLoader(dataset, batch_size=batch_size, shuffle=False)
|
|
19
|
+
|
|
20
|
+
def consume() -> int:
|
|
21
|
+
count = 0
|
|
22
|
+
for xb, _ in loader:
|
|
23
|
+
count += xb.shape[0]
|
|
24
|
+
return count
|
|
25
|
+
|
|
26
|
+
for _ in range(warmup):
|
|
27
|
+
consume()
|
|
28
|
+
start = time.perf_counter()
|
|
29
|
+
for _ in range(iterations):
|
|
30
|
+
consume()
|
|
31
|
+
elapsed = (time.perf_counter() - start) / iterations
|
|
32
|
+
|
|
33
|
+
print(f"Python: {platform.python_version()}")
|
|
34
|
+
print(f"NumPy: {np.__version__}")
|
|
35
|
+
print(f"ModelStudio: {ms.__version__}")
|
|
36
|
+
print(f"Samples: {samples}")
|
|
37
|
+
print(f"Features: {features}")
|
|
38
|
+
print(f"Batch size: {batch_size}")
|
|
39
|
+
print(f"Warmup: {warmup}")
|
|
40
|
+
print(f"Iterations: {iterations}")
|
|
41
|
+
print("Backend: CPU MVP")
|
|
42
|
+
print(f"DataLoader avg: {elapsed * 1_000:.3f} ms")
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
if __name__ == "__main__":
|
|
46
|
+
main()
|
|
@@ -0,0 +1,37 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import platform
|
|
4
|
+
import time
|
|
5
|
+
|
|
6
|
+
import modelstudio as ms
|
|
7
|
+
import numpy as np
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
def timeit(fn, iterations: int, warmup: int) -> float:
|
|
11
|
+
for _ in range(warmup):
|
|
12
|
+
fn()
|
|
13
|
+
start = time.perf_counter()
|
|
14
|
+
for _ in range(iterations):
|
|
15
|
+
fn()
|
|
16
|
+
return (time.perf_counter() - start) / iterations
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
def main() -> None:
|
|
20
|
+
warmup = 5
|
|
21
|
+
iterations = 50
|
|
22
|
+
shape = (512, 512)
|
|
23
|
+
x = ms.randn(shape)
|
|
24
|
+
|
|
25
|
+
print(f"Python: {platform.python_version()}")
|
|
26
|
+
print(f"NumPy: {np.__version__}")
|
|
27
|
+
print(f"ModelStudio: {ms.__version__}")
|
|
28
|
+
print(f"Operation: dropout shape={shape} p=0.5")
|
|
29
|
+
print(f"Warmup: {warmup}")
|
|
30
|
+
print(f"Iterations: {iterations}")
|
|
31
|
+
print("Backend: CPU only")
|
|
32
|
+
print(f"Dropout avg: {timeit(lambda: ms.dropout(x, p=0.5), iterations, warmup) * 1_000:.3f} ms")
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
if __name__ == "__main__":
|
|
36
|
+
main()
|
|
37
|
+
|
|
@@ -2,6 +2,8 @@ add_library(modelstudio_native STATIC
|
|
|
2
2
|
backends/cpu/cpu_backend.cpp
|
|
3
3
|
backends/cpu/kernels/add.cpp
|
|
4
4
|
backends/cpu/kernels/matmul.cpp
|
|
5
|
+
backends/cpu/kernels/mul.cpp
|
|
6
|
+
backends/cpu/kernels/relu.cpp
|
|
5
7
|
)
|
|
6
8
|
|
|
7
9
|
target_include_directories(modelstudio_native PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})
|
|
@@ -10,8 +10,16 @@ Tensor CPUBackend::add(const Tensor& lhs, const Tensor& rhs) {
|
|
|
10
10
|
return add_kernel(lhs, rhs);
|
|
11
11
|
}
|
|
12
12
|
|
|
13
|
+
Tensor CPUBackend::mul(const Tensor& lhs, const Tensor& rhs) {
|
|
14
|
+
return mul_kernel(lhs, rhs);
|
|
15
|
+
}
|
|
16
|
+
|
|
13
17
|
Tensor CPUBackend::matmul(const Tensor& lhs, const Tensor& rhs) {
|
|
14
18
|
return matmul_kernel(lhs, rhs);
|
|
15
19
|
}
|
|
16
20
|
|
|
21
|
+
Tensor CPUBackend::relu(const Tensor& input) {
|
|
22
|
+
return relu_kernel(input);
|
|
23
|
+
}
|
|
24
|
+
|
|
17
25
|
} // namespace modelstudio::cpu
|
|
@@ -5,14 +5,18 @@
|
|
|
5
5
|
namespace modelstudio::cpu {
|
|
6
6
|
|
|
7
7
|
Tensor add_kernel(const Tensor& lhs, const Tensor& rhs);
|
|
8
|
+
Tensor mul_kernel(const Tensor& lhs, const Tensor& rhs);
|
|
8
9
|
Tensor matmul_kernel(const Tensor& lhs, const Tensor& rhs);
|
|
10
|
+
Tensor relu_kernel(const Tensor& input);
|
|
9
11
|
|
|
10
12
|
class CPUBackend final : public Backend {
|
|
11
13
|
public:
|
|
12
14
|
std::string_view name() const override { return "cpu"; }
|
|
13
15
|
Tensor empty(const Shape& shape, DType dtype) override;
|
|
14
16
|
Tensor add(const Tensor& lhs, const Tensor& rhs) override;
|
|
17
|
+
Tensor mul(const Tensor& lhs, const Tensor& rhs) override;
|
|
15
18
|
Tensor matmul(const Tensor& lhs, const Tensor& rhs) override;
|
|
19
|
+
Tensor relu(const Tensor& input) override;
|
|
16
20
|
};
|
|
17
21
|
|
|
18
22
|
} // namespace modelstudio::cpu
|
|
@@ -0,0 +1,43 @@
|
|
|
1
|
+
#include "backends/cpu/cpu_backend.hpp"
|
|
2
|
+
|
|
3
|
+
namespace modelstudio::cpu {
|
|
4
|
+
namespace {
|
|
5
|
+
|
|
6
|
+
template <typename T>
|
|
7
|
+
void mul_typed(const Tensor& lhs, const Tensor& rhs, Tensor& out) {
|
|
8
|
+
const auto* lhs_ptr = static_cast<const T*>(lhs.data());
|
|
9
|
+
const auto* rhs_ptr = static_cast<const T*>(rhs.data());
|
|
10
|
+
auto* out_ptr = static_cast<T*>(out.data());
|
|
11
|
+
for (std::int64_t i = 0; i < lhs.numel(); ++i) {
|
|
12
|
+
out_ptr[i] = lhs_ptr[i] * rhs_ptr[i];
|
|
13
|
+
}
|
|
14
|
+
}
|
|
15
|
+
|
|
16
|
+
} // namespace
|
|
17
|
+
|
|
18
|
+
Tensor mul_kernel(const Tensor& lhs, const Tensor& rhs) {
|
|
19
|
+
if (lhs.shape() != rhs.shape() || lhs.dtype() != rhs.dtype()) {
|
|
20
|
+
throw Error("native CPU mul currently requires identical shape and dtype");
|
|
21
|
+
}
|
|
22
|
+
Tensor out(lhs.shape(), lhs.dtype(), lhs.device());
|
|
23
|
+
switch (lhs.dtype()) {
|
|
24
|
+
case DType::Float32:
|
|
25
|
+
mul_typed<float>(lhs, rhs, out);
|
|
26
|
+
break;
|
|
27
|
+
case DType::Float64:
|
|
28
|
+
mul_typed<double>(lhs, rhs, out);
|
|
29
|
+
break;
|
|
30
|
+
case DType::Int32:
|
|
31
|
+
mul_typed<std::int32_t>(lhs, rhs, out);
|
|
32
|
+
break;
|
|
33
|
+
case DType::Int64:
|
|
34
|
+
mul_typed<std::int64_t>(lhs, rhs, out);
|
|
35
|
+
break;
|
|
36
|
+
case DType::Bool:
|
|
37
|
+
mul_typed<bool>(lhs, rhs, out);
|
|
38
|
+
break;
|
|
39
|
+
}
|
|
40
|
+
return out;
|
|
41
|
+
}
|
|
42
|
+
|
|
43
|
+
} // namespace modelstudio::cpu
|