openarchx 0.1.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.
Files changed (60) hide show
  1. openarchx-0.1.0/LICENSE +21 -0
  2. openarchx-0.1.0/MANIFEST.in +5 -0
  3. openarchx-0.1.0/PKG-INFO +180 -0
  4. openarchx-0.1.0/README.md +147 -0
  5. openarchx-0.1.0/examples/core_test.py +147 -0
  6. openarchx-0.1.0/examples/cuda_test.py +61 -0
  7. openarchx-0.1.0/examples/hybrid_sequence_classifier.py +200 -0
  8. openarchx-0.1.0/examples/integration_example.py +212 -0
  9. openarchx-0.1.0/examples/mnist_cnn.py +210 -0
  10. openarchx-0.1.0/examples/mnist_example.py +109 -0
  11. openarchx-0.1.0/examples/model_saving/README.md +85 -0
  12. openarchx-0.1.0/examples/model_saving/model_saving_example.py +249 -0
  13. openarchx-0.1.0/examples/quantum_transformer_example.py +105 -0
  14. openarchx-0.1.0/examples/simple_transformer_test.py +68 -0
  15. openarchx-0.1.0/openarchx/__init__.py +11 -0
  16. openarchx-0.1.0/openarchx/core/tensor.py +179 -0
  17. openarchx-0.1.0/openarchx/cuda/__init__.py +27 -0
  18. openarchx-0.1.0/openarchx/cuda/cuda_ops.py +296 -0
  19. openarchx-0.1.0/openarchx/layers/activations.py +63 -0
  20. openarchx-0.1.0/openarchx/layers/base.py +40 -0
  21. openarchx-0.1.0/openarchx/layers/cnn.py +145 -0
  22. openarchx-0.1.0/openarchx/layers/transformer.py +131 -0
  23. openarchx-0.1.0/openarchx/nn/__init__.py +26 -0
  24. openarchx-0.1.0/openarchx/nn/activations.py +127 -0
  25. openarchx-0.1.0/openarchx/nn/containers.py +174 -0
  26. openarchx-0.1.0/openarchx/nn/dropout.py +121 -0
  27. openarchx-0.1.0/openarchx/nn/layers.py +338 -0
  28. openarchx-0.1.0/openarchx/nn/losses.py +156 -0
  29. openarchx-0.1.0/openarchx/nn/module.py +18 -0
  30. openarchx-0.1.0/openarchx/nn/padding.py +120 -0
  31. openarchx-0.1.0/openarchx/nn/pooling.py +318 -0
  32. openarchx-0.1.0/openarchx/nn/rnn.py +226 -0
  33. openarchx-0.1.0/openarchx/nn/transformers.py +187 -0
  34. openarchx-0.1.0/openarchx/optimizers/adam.py +49 -0
  35. openarchx-0.1.0/openarchx/optimizers/adaptive.py +63 -0
  36. openarchx-0.1.0/openarchx/optimizers/base.py +24 -0
  37. openarchx-0.1.0/openarchx/optimizers/modern.py +98 -0
  38. openarchx-0.1.0/openarchx/optimizers/optx.py +91 -0
  39. openarchx-0.1.0/openarchx/optimizers/sgd.py +63 -0
  40. openarchx-0.1.0/openarchx/quantum/circuit.py +92 -0
  41. openarchx-0.1.0/openarchx/quantum/gates.py +126 -0
  42. openarchx-0.1.0/openarchx/utils/__init__.py +50 -0
  43. openarchx-0.1.0/openarchx/utils/data.py +229 -0
  44. openarchx-0.1.0/openarchx/utils/huggingface.py +288 -0
  45. openarchx-0.1.0/openarchx/utils/losses.py +21 -0
  46. openarchx-0.1.0/openarchx/utils/model_io.py +553 -0
  47. openarchx-0.1.0/openarchx/utils/pytorch.py +420 -0
  48. openarchx-0.1.0/openarchx/utils/tensorflow.py +467 -0
  49. openarchx-0.1.0/openarchx/utils/transforms.py +259 -0
  50. openarchx-0.1.0/openarchx.egg-info/PKG-INFO +180 -0
  51. openarchx-0.1.0/openarchx.egg-info/SOURCES.txt +58 -0
  52. openarchx-0.1.0/openarchx.egg-info/dependency_links.txt +1 -0
  53. openarchx-0.1.0/openarchx.egg-info/requires.txt +17 -0
  54. openarchx-0.1.0/openarchx.egg-info/top_level.txt +2 -0
  55. openarchx-0.1.0/pyproject.toml +49 -0
  56. openarchx-0.1.0/setup.cfg +4 -0
  57. openarchx-0.1.0/setup.py +54 -0
  58. openarchx-0.1.0/tests/__init__.py +1 -0
  59. openarchx-0.1.0/tests/test_cuda_ops.py +205 -0
  60. openarchx-0.1.0/tests/test_integrations.py +236 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2023 OpenArchX
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,5 @@
1
+ include LICENSE
2
+ include README.md
3
+ include MANIFEST.in
4
+ recursive-include examples *.py
5
+ recursive-include examples *.md
@@ -0,0 +1,180 @@
1
+ Metadata-Version: 2.4
2
+ Name: openarchx
3
+ Version: 0.1.0
4
+ Summary: A lightweight and extensible deep learning framework with native model serialization
5
+ Home-page: https://github.com/openarchx/openarchx
6
+ Author: OpenArchX Team
7
+ Author-email: OpenArchX Team <info@openarchx.org>
8
+ License-Expression: MIT
9
+ Project-URL: Homepage, https://github.com/openarchx/openarchx
10
+ Project-URL: Bug Tracker, https://github.com/openarchx/openarchx/issues
11
+ Classifier: Programming Language :: Python :: 3
12
+ Classifier: Operating System :: OS Independent
13
+ Requires-Python: >=3.7
14
+ Description-Content-Type: text/markdown
15
+ License-File: LICENSE
16
+ Requires-Dist: numpy>=1.19.0
17
+ Provides-Extra: pytorch
18
+ Requires-Dist: torch>=1.7.0; extra == "pytorch"
19
+ Provides-Extra: tensorflow
20
+ Requires-Dist: tensorflow>=2.4.0; extra == "tensorflow"
21
+ Provides-Extra: huggingface
22
+ Requires-Dist: transformers>=4.0.0; extra == "huggingface"
23
+ Requires-Dist: datasets>=1.0.0; extra == "huggingface"
24
+ Provides-Extra: all
25
+ Requires-Dist: torch>=1.7.0; extra == "all"
26
+ Requires-Dist: tensorflow>=2.4.0; extra == "all"
27
+ Requires-Dist: transformers>=4.0.0; extra == "all"
28
+ Requires-Dist: datasets>=1.0.0; extra == "all"
29
+ Dynamic: author
30
+ Dynamic: home-page
31
+ Dynamic: license-file
32
+ Dynamic: requires-python
33
+
34
+ # OpenArchX
35
+
36
+ [![PyPI version](https://badge.fury.io/py/openarchx.svg)](https://badge.fury.io/py/openarchx)
37
+ [![GitHub license](https://img.shields.io/github/license/openarchx/openarchx)](https://github.com/openarchx/openarchx/blob/main/LICENSE)
38
+ [![Python Versions](https://img.shields.io/pypi/pyversions/openarchx.svg)](https://pypi.org/project/openarchx/)
39
+
40
+ A lightweight and extensible deep learning framework in pure Python with native model serialization support.
41
+
42
+ ## Features
43
+
44
+ - Simple and clean API inspired by modern deep learning frameworks
45
+ - Native `.oaxm` model serialization format
46
+ - Seamless integration with PyTorch, TensorFlow, and Hugging Face
47
+ - Framework-agnostic design for maximum flexibility
48
+ - Pure Python implementation with minimal dependencies
49
+
50
+ ## Installation
51
+
52
+ ### Basic Installation
53
+
54
+ ```bash
55
+ pip install openarchx
56
+ ```
57
+
58
+ ### With Framework Integration Support
59
+
60
+ ```bash
61
+ # For PyTorch integration
62
+ pip install openarchx[pytorch]
63
+
64
+ # For TensorFlow integration
65
+ pip install openarchx[tensorflow]
66
+
67
+ # For Hugging Face integration
68
+ pip install openarchx[huggingface]
69
+
70
+ # For all integrations
71
+ pip install openarchx[all]
72
+ ```
73
+
74
+ ## Quick Start
75
+
76
+ ```python
77
+ import numpy as np
78
+ import openarchx as ox
79
+ from openarchx.nn import Sequential, Dense, ReLU
80
+ from openarchx.core import Tensor
81
+ from openarchx.utils import save_model, load_model
82
+
83
+ # Create a model
84
+ model = Sequential([
85
+ Dense(10, input_dim=5),
86
+ ReLU(),
87
+ Dense(1)
88
+ ])
89
+
90
+ # Generate dummy data
91
+ X = np.random.randn(100, 5).astype(np.float32)
92
+ y = np.sum(X * np.array([0.2, 0.5, -0.3, 0.7, -0.1]), axis=1, keepdims=True)
93
+ X, y = Tensor(X), Tensor(y)
94
+
95
+ # Train the model
96
+ optimizer = ox.optim.SGD(model.parameters(), learning_rate=0.01)
97
+ loss_fn = ox.losses.MSELoss()
98
+
99
+ for epoch in range(10):
100
+ # Forward pass
101
+ y_pred = model(X)
102
+ loss = loss_fn(y_pred, y)
103
+
104
+ # Backward pass
105
+ optimizer.zero_grad()
106
+ loss.backward()
107
+ optimizer.step()
108
+
109
+ print(f"Epoch {epoch}: Loss = {loss.data}")
110
+
111
+ # Save the model to .oaxm format
112
+ save_model(model, "my_model.oaxm")
113
+
114
+ # Load the model
115
+ loaded_model = load_model("my_model.oaxm", model_class=Sequential)
116
+ ```
117
+
118
+ ## Model Serialization with .oaxm
119
+
120
+ OpenArchX provides a native model serialization format called `.oaxm` (OpenArchX Model):
121
+
122
+ ```python
123
+ # Save a model with metadata
124
+ metadata = {
125
+ "description": "My trained model",
126
+ "version": "1.0.0",
127
+ "author": "Your Name"
128
+ }
129
+ save_model(model, "model.oaxm", metadata=metadata)
130
+
131
+ # Convert from PyTorch
132
+ from openarchx.utils import convert_from_pytorch
133
+ convert_from_pytorch(torch_model, "converted_model.oaxm")
134
+
135
+ # Convert from TensorFlow
136
+ from openarchx.utils import convert_from_tensorflow
137
+ convert_from_tensorflow(tf_model, "converted_model.oaxm")
138
+ ```
139
+
140
+ ## Framework Integration
141
+
142
+ ### PyTorch Integration
143
+
144
+ ```python
145
+ import torch
146
+ import torch.nn as nn
147
+ from openarchx.utils import get_pytorch_model_adapter
148
+
149
+ # Convert PyTorch model to OpenArchX
150
+ pt_model = nn.Sequential(
151
+ nn.Linear(10, 5),
152
+ nn.ReLU(),
153
+ nn.Linear(5, 1)
154
+ )
155
+
156
+ # Create an adapter to use the PyTorch model in OpenArchX
157
+ adapted_model = get_pytorch_model_adapter(pt_model)
158
+ output = adapted_model(Tensor(np.random.randn(1, 10)))
159
+ ```
160
+
161
+ ### TensorFlow Integration
162
+
163
+ ```python
164
+ import tensorflow as tf
165
+ from openarchx.utils import get_tensorflow_model_adapter
166
+
167
+ # Use TensorFlow model in OpenArchX
168
+ tf_model = tf.keras.Sequential([
169
+ tf.keras.layers.Dense(5, activation='relu', input_shape=(10,)),
170
+ tf.keras.layers.Dense(1)
171
+ ])
172
+
173
+ # Create an adapter to use the TensorFlow model in OpenArchX
174
+ adapted_model = get_tensorflow_model_adapter(tf_model)
175
+ output = adapted_model(Tensor(np.random.randn(1, 10)))
176
+ ```
177
+
178
+ ## License
179
+
180
+ This project is licensed under the MIT License - see the LICENSE file for details.
@@ -0,0 +1,147 @@
1
+ # OpenArchX
2
+
3
+ [![PyPI version](https://badge.fury.io/py/openarchx.svg)](https://badge.fury.io/py/openarchx)
4
+ [![GitHub license](https://img.shields.io/github/license/openarchx/openarchx)](https://github.com/openarchx/openarchx/blob/main/LICENSE)
5
+ [![Python Versions](https://img.shields.io/pypi/pyversions/openarchx.svg)](https://pypi.org/project/openarchx/)
6
+
7
+ A lightweight and extensible deep learning framework in pure Python with native model serialization support.
8
+
9
+ ## Features
10
+
11
+ - Simple and clean API inspired by modern deep learning frameworks
12
+ - Native `.oaxm` model serialization format
13
+ - Seamless integration with PyTorch, TensorFlow, and Hugging Face
14
+ - Framework-agnostic design for maximum flexibility
15
+ - Pure Python implementation with minimal dependencies
16
+
17
+ ## Installation
18
+
19
+ ### Basic Installation
20
+
21
+ ```bash
22
+ pip install openarchx
23
+ ```
24
+
25
+ ### With Framework Integration Support
26
+
27
+ ```bash
28
+ # For PyTorch integration
29
+ pip install openarchx[pytorch]
30
+
31
+ # For TensorFlow integration
32
+ pip install openarchx[tensorflow]
33
+
34
+ # For Hugging Face integration
35
+ pip install openarchx[huggingface]
36
+
37
+ # For all integrations
38
+ pip install openarchx[all]
39
+ ```
40
+
41
+ ## Quick Start
42
+
43
+ ```python
44
+ import numpy as np
45
+ import openarchx as ox
46
+ from openarchx.nn import Sequential, Dense, ReLU
47
+ from openarchx.core import Tensor
48
+ from openarchx.utils import save_model, load_model
49
+
50
+ # Create a model
51
+ model = Sequential([
52
+ Dense(10, input_dim=5),
53
+ ReLU(),
54
+ Dense(1)
55
+ ])
56
+
57
+ # Generate dummy data
58
+ X = np.random.randn(100, 5).astype(np.float32)
59
+ y = np.sum(X * np.array([0.2, 0.5, -0.3, 0.7, -0.1]), axis=1, keepdims=True)
60
+ X, y = Tensor(X), Tensor(y)
61
+
62
+ # Train the model
63
+ optimizer = ox.optim.SGD(model.parameters(), learning_rate=0.01)
64
+ loss_fn = ox.losses.MSELoss()
65
+
66
+ for epoch in range(10):
67
+ # Forward pass
68
+ y_pred = model(X)
69
+ loss = loss_fn(y_pred, y)
70
+
71
+ # Backward pass
72
+ optimizer.zero_grad()
73
+ loss.backward()
74
+ optimizer.step()
75
+
76
+ print(f"Epoch {epoch}: Loss = {loss.data}")
77
+
78
+ # Save the model to .oaxm format
79
+ save_model(model, "my_model.oaxm")
80
+
81
+ # Load the model
82
+ loaded_model = load_model("my_model.oaxm", model_class=Sequential)
83
+ ```
84
+
85
+ ## Model Serialization with .oaxm
86
+
87
+ OpenArchX provides a native model serialization format called `.oaxm` (OpenArchX Model):
88
+
89
+ ```python
90
+ # Save a model with metadata
91
+ metadata = {
92
+ "description": "My trained model",
93
+ "version": "1.0.0",
94
+ "author": "Your Name"
95
+ }
96
+ save_model(model, "model.oaxm", metadata=metadata)
97
+
98
+ # Convert from PyTorch
99
+ from openarchx.utils import convert_from_pytorch
100
+ convert_from_pytorch(torch_model, "converted_model.oaxm")
101
+
102
+ # Convert from TensorFlow
103
+ from openarchx.utils import convert_from_tensorflow
104
+ convert_from_tensorflow(tf_model, "converted_model.oaxm")
105
+ ```
106
+
107
+ ## Framework Integration
108
+
109
+ ### PyTorch Integration
110
+
111
+ ```python
112
+ import torch
113
+ import torch.nn as nn
114
+ from openarchx.utils import get_pytorch_model_adapter
115
+
116
+ # Convert PyTorch model to OpenArchX
117
+ pt_model = nn.Sequential(
118
+ nn.Linear(10, 5),
119
+ nn.ReLU(),
120
+ nn.Linear(5, 1)
121
+ )
122
+
123
+ # Create an adapter to use the PyTorch model in OpenArchX
124
+ adapted_model = get_pytorch_model_adapter(pt_model)
125
+ output = adapted_model(Tensor(np.random.randn(1, 10)))
126
+ ```
127
+
128
+ ### TensorFlow Integration
129
+
130
+ ```python
131
+ import tensorflow as tf
132
+ from openarchx.utils import get_tensorflow_model_adapter
133
+
134
+ # Use TensorFlow model in OpenArchX
135
+ tf_model = tf.keras.Sequential([
136
+ tf.keras.layers.Dense(5, activation='relu', input_shape=(10,)),
137
+ tf.keras.layers.Dense(1)
138
+ ])
139
+
140
+ # Create an adapter to use the TensorFlow model in OpenArchX
141
+ adapted_model = get_tensorflow_model_adapter(tf_model)
142
+ output = adapted_model(Tensor(np.random.randn(1, 10)))
143
+ ```
144
+
145
+ ## License
146
+
147
+ This project is licensed under the MIT License - see the LICENSE file for details.
@@ -0,0 +1,147 @@
1
+ import numpy as np
2
+ import traceback
3
+ import sys
4
+ from openarchx.core.tensor import Tensor
5
+ from openarchx.nn.module import Module
6
+ from openarchx.layers.cnn import Conv2d, MaxPool2d
7
+ from openarchx.layers.base import Linear
8
+ from openarchx.optimizers.adam import Adam
9
+ from openarchx.layers.activations import relu
10
+
11
+ def print_separator():
12
+ print("\n" + "="*50 + "\n")
13
+
14
+ def softmax(x):
15
+ """Stable softmax implementation"""
16
+ if isinstance(x, Tensor):
17
+ x = x.data
18
+ x = x - np.max(x, axis=1, keepdims=True) # For numerical stability
19
+ exp_x = np.exp(x)
20
+ return exp_x / exp_x.sum(axis=1, keepdims=True)
21
+
22
+ def cross_entropy_loss(pred, target):
23
+ """Stable cross entropy loss calculation"""
24
+ if isinstance(pred, Tensor):
25
+ pred = pred.data
26
+ if isinstance(target, Tensor):
27
+ target = target.data
28
+
29
+ epsilon = 1e-7
30
+ pred_clipped = np.clip(pred, epsilon, 1 - epsilon)
31
+ return -np.mean(np.sum(target * np.log(pred_clipped), axis=1))
32
+
33
+ def test_tensor_ops():
34
+ print_separator()
35
+ print("Testing basic tensor operations...")
36
+
37
+ try:
38
+ # Test creation and basic ops
39
+ a = Tensor(np.array([[1, 2], [3, 4]]), requires_grad=True)
40
+ b = Tensor(np.array([[5, 6], [7, 8]]), requires_grad=True)
41
+
42
+ print("Addition:")
43
+ c = a + b
44
+ print(c.data)
45
+
46
+ print("\nMatrix multiplication:")
47
+ d = a @ b
48
+ print(d.data)
49
+
50
+ print("\nReshape and transpose:")
51
+ e = a.reshape(4)
52
+ print(f"Reshaped: {e.data}")
53
+ f = a.transpose(1, 0)
54
+ print(f"Transposed: {f.data}")
55
+
56
+ print("\nTensor operations test passed!")
57
+ except Exception as e:
58
+ print(f"Tensor operations error: {e}")
59
+ traceback.print_exc()
60
+
61
+ def test_conv_layer():
62
+ print_separator()
63
+ print("Testing convolution layer...")
64
+
65
+ try:
66
+ # Create small input
67
+ x = np.random.randn(2, 3, 4, 4) # [batch_size, channels, height, width]
68
+ x_tensor = Tensor(x)
69
+
70
+ # Create conv layer
71
+ conv = Conv2d(in_channels=3, out_channels=2, kernel_size=3, padding=1)
72
+
73
+ # Forward pass
74
+ out = conv(x_tensor)
75
+ print(f"Conv2d input shape: {x.shape}")
76
+ print(f"Conv2d output shape: {out.data.shape}")
77
+ print("\nConvolution layer test passed!")
78
+
79
+ except Exception as e:
80
+ print(f"Conv2d error: {e}")
81
+ traceback.print_exc()
82
+
83
+ class SimpleNet(Module):
84
+ def __init__(self):
85
+ super().__init__()
86
+ self.conv1 = Conv2d(1, 4, kernel_size=3, padding=1)
87
+ self.fc1 = Linear(64, 10) # Fixed: 4 * 4 * 4 = 64 input features
88
+
89
+ def forward(self, x):
90
+ x = relu(self.conv1(x))
91
+ batch_size = x.data.shape[0]
92
+ x = x.reshape(batch_size, -1) # Flatten conv output
93
+ x = self.fc1(x)
94
+ probs = softmax(x) # Apply softmax to logits
95
+ return Tensor(probs, requires_grad=True)
96
+
97
+ def test_small_network():
98
+ print_separator()
99
+ print("Testing small network...")
100
+
101
+ try:
102
+ # Create small input
103
+ np.random.seed(42) # For reproducibility
104
+ x = np.random.randn(2, 1, 4, 4) # [batch_size, channels, height, width]
105
+ y = np.zeros((2, 10)) # One-hot encoded targets
106
+ y[:, 0] = 1 # Set first class as target
107
+
108
+ # Create model and optimizer
109
+ model = SimpleNet()
110
+ optimizer = Adam(model.parameters(), lr=0.01)
111
+
112
+ # Training step
113
+ x_tensor = Tensor(x)
114
+ y_tensor = Tensor(y)
115
+
116
+ print(f"Input shape: {x.shape}")
117
+
118
+ # Forward pass
119
+ pred = model(x_tensor)
120
+ print(f"Output shape: {pred.data.shape}")
121
+
122
+ # Calculate loss
123
+ loss = cross_entropy_loss(pred, y_tensor)
124
+ print(f"Initial loss: {loss:.4f}")
125
+
126
+ # Backward pass
127
+ loss_tensor = Tensor(np.array(loss), requires_grad=True)
128
+ optimizer.zero_grad()
129
+ loss_tensor.backward()
130
+ optimizer.step()
131
+
132
+ # Another forward pass to check if loss decreased
133
+ pred = model(x_tensor)
134
+ new_loss = cross_entropy_loss(pred, y_tensor)
135
+ print(f"Loss after one step: {new_loss:.4f}")
136
+ print("\nNetwork training test passed!")
137
+
138
+ except Exception as e:
139
+ print(f"Training error: {e}")
140
+ traceback.print_exc()
141
+
142
+ if __name__ == "__main__":
143
+ print("Starting core functionality tests...")
144
+ test_tensor_ops()
145
+ test_conv_layer()
146
+ test_small_network()
147
+ print("\nTests completed!")
@@ -0,0 +1,61 @@
1
+ from openarchx.core.tensor import Tensor
2
+ import numpy as np
3
+ import time
4
+
5
+ def benchmark_matmul(size=1000, device='cpu'):
6
+ """Benchmark matrix multiplication on CPU vs GPU"""
7
+ # Create random matrices
8
+ a = np.random.randn(size, size).astype(np.float32)
9
+ b = np.random.randn(size, size).astype(np.float32)
10
+
11
+ # Convert to tensors
12
+ ta = Tensor(a, device=device)
13
+ tb = Tensor(b, device=device)
14
+
15
+ # Warmup
16
+ _ = ta @ tb
17
+
18
+ # Benchmark
19
+ start = time.time()
20
+ for _ in range(10):
21
+ c = ta @ tb
22
+ end = time.time()
23
+
24
+ return (end - start) / 10
25
+
26
+ def test_cuda_ops():
27
+ """Test basic CUDA operations"""
28
+ try:
29
+ # Test tensor creation and device movement
30
+ a = Tensor(np.array([[1, 2], [3, 4]]), device='cuda')
31
+ b = Tensor(np.array([[5, 6], [7, 8]]), device='cuda')
32
+
33
+ print("CUDA Tensor Operations Test:")
34
+ print("----------------------------")
35
+
36
+ # Test addition
37
+ c = a + b
38
+ print(f"Addition result:\n{c}")
39
+
40
+ # Test matrix multiplication
41
+ d = a @ b
42
+ print(f"\nMatrix multiplication result:\n{d}")
43
+
44
+ # Test element-wise multiplication
45
+ e = a * b
46
+ print(f"\nElement-wise multiplication result:\n{e}")
47
+
48
+ # Benchmark large matrix multiplication
49
+ print("\nBenchmarking matrix multiplication:")
50
+ cpu_time = benchmark_matmul(size=1000, device='cpu')
51
+ gpu_time = benchmark_matmul(size=1000, device='cuda')
52
+
53
+ print(f"CPU average time: {cpu_time:.4f} seconds")
54
+ print(f"GPU average time: {gpu_time:.4f} seconds")
55
+ print(f"Speedup: {cpu_time/gpu_time:.2f}x")
56
+
57
+ except Exception as e:
58
+ print(f"Error during CUDA test: {str(e)}")
59
+
60
+ if __name__ == "__main__":
61
+ test_cuda_ops()