Skip to content

Commit

Permalink
refactor: array_library.py, update docstrings, add license
Browse files Browse the repository at this point in the history
  • Loading branch information
sradc committed Mar 23, 2022
1 parent d17ad9d commit da70489
Showing 1 changed file with 20 additions and 9 deletions.
29 changes: 20 additions & 9 deletions smallpebble/array_library.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,21 @@
"""Allows SmallPebble to dynamically switch between NumPy/CuPy,
for its ndarray computations. The default is NumPy.
# Copyright 2021 Sidney Radcliffe
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""Allows SmallPebble to dynamically switch between using NumPy/CuPy,
as its ndarray library. The default is NumPy.
Note to SmallPebble devs:
To SmallPebble devs:
Watch out for cases where NumPy and CuPy differ,
e.g. np.add.at is cupy.scatter_add.
"""
Expand Down Expand Up @@ -32,7 +46,6 @@ def use(array_library: ModuleType) -> None:
sp.use(numpy)
```
"""

global library
library = array_library

Expand All @@ -41,19 +54,17 @@ def __getattr__(name: str):
"""Make this module act as a proxy, for NumPy/CuPy.
Here's an example:
```python
import smallpebble.array_library as array_library
x = array_library.array([1, 2, 3]) ** 2 # <- a NumPy array
```
In this example, a NumPy array is created,
because `array_library.array` results in this function
being called, which then calls `getattr(numpy, "array")`,
which is NumPy's function for creating arrays.
(The above assumes that `library == numpy`, which is the
(The above assumes that `library == numpy`, which is the
default. The same thing would happen but with CuPy arrays,
if `library == cupy`.)
if `library == cupy`.)
"""
return getattr(library, name)

0 comments on commit da70489

Please sign in to comment.