Author: Maciej Fijalkowski <[email protected]>
Branch: numpy-multidim-shards
Changeset: r49546:21a78357a235
Date: 2011-11-19 10:24 +0200
http://bitbucket.org/pypy/pypy/changeset/21a78357a235/
Log: start working on broadcasting - a helper function
diff --git a/pypy/module/micronumpy/interp_numarray.py
b/pypy/module/micronumpy/interp_numarray.py
--- a/pypy/module/micronumpy/interp_numarray.py
+++ b/pypy/module/micronumpy/interp_numarray.py
@@ -39,6 +39,39 @@
shape.append(size)
batch = new_batch
+def shape_agreement(space, shape1, shape2):
+ """ Checks agreement about two shapes with respect to broadcasting. Returns
+ the resulting shape.
+ """
+ lshift = 0
+ rshift = 0
+ if len(shape1) > len(shape2):
+ m = len(shape1)
+ n = len(shape2)
+ rshift = len(shape2) - len(shape1)
+ remainder = shape1
+ else:
+ m = len(shape2)
+ n = len(shape1)
+ lshift = len(shape1) - len(shape2)
+ remainder = shape2
+ endshape = [0] * m
+ for i in range(m - 1, m - n - 1, -1):
+ left = shape1[i + lshift]
+ right = shape2[i + rshift]
+ if left == right:
+ endshape[i] = left
+ elif left == 1:
+ endshape[i] = right
+ elif right == 1:
+ endshape[i] = left
+ else:
+ raise OperationError(space.w_ValueError, space.wrap(
+ "frames are not aligned"))
+ for i in range(m - n):
+ endshape[i] = remainder[i]
+ return endshape
+
def descr_new_array(space, w_subtype, w_item_or_iterable, w_dtype=None,
w_order=NoneNotWrapped):
# find scalar
diff --git a/pypy/module/micronumpy/test/test_numarray.py
b/pypy/module/micronumpy/test/test_numarray.py
--- a/pypy/module/micronumpy/test/test_numarray.py
+++ b/pypy/module/micronumpy/test/test_numarray.py
@@ -1,6 +1,9 @@
+
+import py
from pypy.module.micronumpy.test.test_base import BaseNumpyAppTest
-from pypy.module.micronumpy.interp_numarray import NDimArray
+from pypy.module.micronumpy.interp_numarray import NDimArray, shape_agreement
from pypy.module.micronumpy import signature
+from pypy.interpreter.error import OperationError
from pypy.conftest import gettestobjspace
class MockDtype(object):
@@ -142,6 +145,14 @@
r = s._index_of_single_item(self.space, self.newtuple(1, 1))
assert r == a._index_of_single_item(self.space, self.newtuple(1, 2, 1))
+ def test_shape_agreement(self):
+ assert shape_agreement(self.space, [3], [3]) == [3]
+ assert shape_agreement(self.space, [1, 2, 3], [1, 2, 3]) == [1, 2, 3]
+ py.test.raises(OperationError, shape_agreement, self.space, [2], [3])
+ assert shape_agreement(self.space, [4, 4], []) == [4, 4]
+ assert shape_agreement(self.space, [8, 1, 6, 1], [7, 1, 5]) == [8, 7,
6, 5]
+ assert shape_agreement(self.space, [5, 2], [4, 3, 5, 2]) == [4, 3, 5,
2]
+
class AppTestNumArray(BaseNumpyAppTest):
def test_type(self):
from numpy import array
_______________________________________________
pypy-commit mailing list
[email protected]
http://mail.python.org/mailman/listinfo/pypy-commit
