提交 18e6c04d authored 作者: Brendan Murphy's avatar Brendan Murphy 提交者: Ricardo Vieira

Replace use of `np.MAXDIMS`

`np.MAXDIMS` was removed from the public API and no replacement is given in the migration docs. In numpy <= 1.26, the value of `np.MAXDIMS` was 32. This was often used as a flag to mean `axis=None`. In numpy >= 2.0, the maximum number of dims of an array has been increased to 64; simultaneously, a constant `NPY_RAVEL_AXIS` was added to the C-API to indicate that `axis=None`. In most cases, the use of `np.MAXDIMS` to check for `axis=None` can be replaced by the new constant `NPY_RAVEL_AXIS`. To make this constant accessible when using numpy <= 1.26, I added a function to insert `npy_2_compat.h` into the support code for the affected ops.
上级 46a235ab
...@@ -46,10 +46,21 @@ else: ...@@ -46,10 +46,21 @@ else:
ndarray_c_version = np.core._multiarray_umath._get_ndarray_c_version() # type: ignore[attr-defined] ndarray_c_version = np.core._multiarray_umath._get_ndarray_c_version() # type: ignore[attr-defined]
# used in tests: the type of error thrown if a value is too large for the specified
# numpy data type is different in numpy 2.x
UintOverflowError = OverflowError if using_numpy_2 else TypeError
# to patch up some of the C code, we need to use these special values...
if using_numpy_2: if using_numpy_2:
UintOverflowError = OverflowError numpy_axis_is_none_flag = np.iinfo(np.int32).min # the value of "NPY_RAVEL_AXIS"
else: else:
UintOverflowError = TypeError # 32 is the value used to mark axis = None in Numpy C-API prior to version 2.0
numpy_axis_is_none_flag = 32
# max number of dims is 64 in numpy 2.x; 32 in older versions
numpy_maxdims = 64 if using_numpy_2 else 32
def npy_2_compat_header() -> str: def npy_2_compat_header() -> str:
......
...@@ -2,7 +2,6 @@ import warnings ...@@ -2,7 +2,6 @@ import warnings
from collections.abc import Collection, Iterable from collections.abc import Collection, Iterable
import numpy as np import numpy as np
from numpy.exceptions import AxisError
import pytensor import pytensor
import pytensor.scalar.basic as ps import pytensor.scalar.basic as ps
...@@ -19,10 +18,11 @@ from pytensor.link.c.params_type import ParamsType ...@@ -19,10 +18,11 @@ from pytensor.link.c.params_type import ParamsType
from pytensor.link.c.type import EnumList, Generic from pytensor.link.c.type import EnumList, Generic
from pytensor.npy_2_compat import ( from pytensor.npy_2_compat import (
normalize_axis_index, normalize_axis_index,
normalize_axis_tuple, npy_2_compat_header,
numpy_axis_is_none_flag,
) )
from pytensor.raise_op import Assert from pytensor.raise_op import Assert
from pytensor.scalar import int32 as int_t from pytensor.scalar import int64 as int_t
from pytensor.scalar import upcast from pytensor.scalar import upcast
from pytensor.tensor import TensorLike, as_tensor_variable from pytensor.tensor import TensorLike, as_tensor_variable
from pytensor.tensor import basic as ptb from pytensor.tensor import basic as ptb
...@@ -47,6 +47,7 @@ from pytensor.tensor.math import sum as pt_sum ...@@ -47,6 +47,7 @@ from pytensor.tensor.math import sum as pt_sum
from pytensor.tensor.shape import Shape_i from pytensor.tensor.shape import Shape_i
from pytensor.tensor.subtensor import advanced_inc_subtensor1, set_subtensor from pytensor.tensor.subtensor import advanced_inc_subtensor1, set_subtensor
from pytensor.tensor.type import TensorType, dvector, int_dtypes, integer_dtypes, vector from pytensor.tensor.type import TensorType, dvector, int_dtypes, integer_dtypes, vector
from pytensor.tensor.utils import normalize_reduce_axis
from pytensor.tensor.variable import TensorVariable from pytensor.tensor.variable import TensorVariable
from pytensor.utils import LOCAL_BITWIDTH, PYTHON_INT_BITWIDTH from pytensor.utils import LOCAL_BITWIDTH, PYTHON_INT_BITWIDTH
...@@ -302,7 +303,11 @@ class CumOp(COp): ...@@ -302,7 +303,11 @@ class CumOp(COp):
self.axis = axis self.axis = axis
self.mode = mode self.mode = mode
c_axis = property(lambda self: np.MAXDIMS if self.axis is None else self.axis) @property
def c_axis(self) -> int:
if self.axis is None:
return numpy_axis_is_none_flag
return self.axis
def make_node(self, x): def make_node(self, x):
x = ptb.as_tensor_variable(x) x = ptb.as_tensor_variable(x)
...@@ -359,24 +364,37 @@ class CumOp(COp): ...@@ -359,24 +364,37 @@ class CumOp(COp):
return shapes return shapes
def c_support_code_apply(self, node: Apply, name: str) -> str:
"""Needed to define NPY_RAVEL_AXIS"""
return npy_2_compat_header()
def c_code(self, node, name, inames, onames, sub): def c_code(self, node, name, inames, onames, sub):
(x,) = inames (x,) = inames
(z,) = onames (z,) = onames
fail = sub["fail"] fail = sub["fail"]
params = sub["params"] params = sub["params"]
code = f""" if self.axis is None:
int axis = {params}->c_axis; axis_code = "int axis = NPY_RAVEL_AXIS;\n"
else:
axis_code = f"int axis = {params}->c_axis;\n"
code = (
axis_code
+ f"""
#undef NPY_UF_DBG_TRACING
#define NPY_UF_DBG_TRACING 1
if (axis == 0 && PyArray_NDIM({x}) == 1) if (axis == 0 && PyArray_NDIM({x}) == 1)
axis = NPY_MAXDIMS; axis = NPY_RAVEL_AXIS;
npy_intp shape[1] = {{ PyArray_SIZE({x}) }}; npy_intp shape[1] = {{ PyArray_SIZE({x}) }};
if(axis == NPY_MAXDIMS && !({z} && PyArray_DIMS({z})[0] == shape[0])) if(axis == NPY_RAVEL_AXIS && !({z} && PyArray_DIMS({z})[0] == shape[0]))
{{ {{
Py_XDECREF({z}); Py_XDECREF({z});
{z} = (PyArrayObject*) PyArray_SimpleNew(1, shape, PyArray_TYPE((PyArrayObject*) py_{x})); {z} = (PyArrayObject*) PyArray_SimpleNew(1, shape, PyArray_TYPE({x}));
}} }}
else if(axis != NPY_MAXDIMS && !({z} && PyArray_CompareLists(PyArray_DIMS({z}), PyArray_DIMS({x}), PyArray_NDIM({x})))) else if(axis != NPY_RAVEL_AXIS && !({z} && PyArray_CompareLists(PyArray_DIMS({z}), PyArray_DIMS({x}), PyArray_NDIM({x}))))
{{ {{
Py_XDECREF({z}); Py_XDECREF({z});
{z} = (PyArrayObject*) PyArray_SimpleNew(PyArray_NDIM({x}), PyArray_DIMS({x}), PyArray_TYPE({x})); {z} = (PyArrayObject*) PyArray_SimpleNew(PyArray_NDIM({x}), PyArray_DIMS({x}), PyArray_TYPE({x}));
...@@ -403,11 +421,12 @@ class CumOp(COp): ...@@ -403,11 +421,12 @@ class CumOp(COp):
Py_XDECREF(t); Py_XDECREF(t);
}} }}
""" """
)
return code return code
def c_code_cache_version(self): def c_code_cache_version(self):
return (8,) return (9,)
def __str__(self): def __str__(self):
return f"{self.__class__.__name__}{{{self.axis}, {self.mode}}}" return f"{self.__class__.__name__}{{{self.axis}, {self.mode}}}"
...@@ -598,11 +617,7 @@ def squeeze(x, axis=None): ...@@ -598,11 +617,7 @@ def squeeze(x, axis=None):
elif not isinstance(axis, Collection): elif not isinstance(axis, Collection):
axis = (axis,) axis = (axis,)
# scalar inputs are treated as 1D regarding axis in this `Op` axis = normalize_reduce_axis(axis, ndim=_x.ndim)
try:
axis = normalize_axis_tuple(axis, ndim=max(1, _x.ndim))
except AxisError:
raise AxisError(axis, ndim=_x.ndim)
if not axis: if not axis:
# Nothing to do # Nothing to do
......
...@@ -13,7 +13,11 @@ from pytensor.graph.op import Op ...@@ -13,7 +13,11 @@ from pytensor.graph.op import Op
from pytensor.graph.replace import _vectorize_node from pytensor.graph.replace import _vectorize_node
from pytensor.link.c.op import COp from pytensor.link.c.op import COp
from pytensor.link.c.params_type import ParamsType from pytensor.link.c.params_type import ParamsType
from pytensor.npy_2_compat import normalize_axis_tuple from pytensor.npy_2_compat import (
normalize_axis_tuple,
npy_2_compat_header,
numpy_axis_is_none_flag,
)
from pytensor.printing import pprint from pytensor.printing import pprint
from pytensor.raise_op import Assert from pytensor.raise_op import Assert
from pytensor.scalar.basic import BinaryScalarOp from pytensor.scalar.basic import BinaryScalarOp
...@@ -160,7 +164,7 @@ class Argmax(COp): ...@@ -160,7 +164,7 @@ class Argmax(COp):
c_axis = np.int64(self.axis[0]) c_axis = np.int64(self.axis[0])
else: else:
# The value here doesn't matter, it won't be used # The value here doesn't matter, it won't be used
c_axis = np.int64(-1) c_axis = numpy_axis_is_none_flag
return self.params_type.get_params(c_axis=c_axis) return self.params_type.get_params(c_axis=c_axis)
def make_node(self, x): def make_node(self, x):
...@@ -203,13 +207,17 @@ class Argmax(COp): ...@@ -203,13 +207,17 @@ class Argmax(COp):
max_idx[0] = np.asarray(np.argmax(reshaped_x, axis=-1), dtype="int64") max_idx[0] = np.asarray(np.argmax(reshaped_x, axis=-1), dtype="int64")
def c_support_code_apply(self, node: Apply, name: str) -> str:
"""Needed to define NPY_RAVEL_AXIS"""
return npy_2_compat_header()
def c_code(self, node, name, inp, out, sub): def c_code(self, node, name, inp, out, sub):
(x,) = inp (x,) = inp
(argmax,) = out (argmax,) = out
fail = sub["fail"] fail = sub["fail"]
params = sub["params"] params = sub["params"]
if self.axis is None: if self.axis is None:
axis_code = "axis = NPY_MAXDIMS;" axis_code = "axis = NPY_RAVEL_AXIS;"
else: else:
if len(self.axis) != 1: if len(self.axis) != 1:
raise NotImplementedError() raise NotImplementedError()
......
...@@ -6,6 +6,7 @@ import scipy ...@@ -6,6 +6,7 @@ import scipy
from pytensor.graph.basic import Apply from pytensor.graph.basic import Apply
from pytensor.graph.replace import _vectorize_node from pytensor.graph.replace import _vectorize_node
from pytensor.link.c.op import COp from pytensor.link.c.op import COp
from pytensor.npy_2_compat import npy_2_compat_header
from pytensor.tensor.basic import as_tensor_variable from pytensor.tensor.basic import as_tensor_variable
from pytensor.tensor.elemwise import get_normalized_batch_axes from pytensor.tensor.elemwise import get_normalized_batch_axes
from pytensor.tensor.math import gamma, gammaln, log, neg, sum from pytensor.tensor.math import gamma, gammaln, log, neg, sum
...@@ -60,12 +61,16 @@ class SoftmaxGrad(COp): ...@@ -60,12 +61,16 @@ class SoftmaxGrad(COp):
return [shape[1]] return [shape[1]]
def c_code_cache_version(self): def c_code_cache_version(self):
return (4,) return (5,)
def c_support_code_apply(self, node: Apply, name: str) -> str:
# return super().c_support_code_apply(node, name)
return npy_2_compat_header()
def c_code(self, node, name, inp, out, sub): def c_code(self, node, name, inp, out, sub):
dy, sm = inp dy, sm = inp
(dx,) = out (dx,) = out
axis = self.axis if self.axis is not None else np.MAXDIMS axis = self.axis if self.axis is not None else "NPY_RAVEL_AXIS"
fail = sub["fail"] fail = sub["fail"]
return dedent( return dedent(
...@@ -79,7 +84,7 @@ class SoftmaxGrad(COp): ...@@ -79,7 +84,7 @@ class SoftmaxGrad(COp):
int sm_ndim = PyArray_NDIM({sm}); int sm_ndim = PyArray_NDIM({sm});
int axis = {axis}; int axis = {axis};
int iterate_axis = !(axis == NPY_MAXDIMS || sm_ndim == 1); int iterate_axis = !(axis == NPY_RAVEL_AXIS || sm_ndim == 1);
// Validate inputs // Validate inputs
if ((PyArray_TYPE({dy}) != NPY_DOUBLE) && if ((PyArray_TYPE({dy}) != NPY_DOUBLE) &&
...@@ -95,13 +100,15 @@ class SoftmaxGrad(COp): ...@@ -95,13 +100,15 @@ class SoftmaxGrad(COp):
{fail}; {fail};
}} }}
if (axis < 0) axis = sm_ndim + axis; if (iterate_axis)
if ((axis < 0) || (iterate_axis && (axis > sm_ndim)))
{{ {{
PyErr_SetString(PyExc_ValueError, "invalid axis in SoftmaxGrad"); if (axis < 0) axis = sm_ndim + axis;
{fail}; if ((axis < 0) || (iterate_axis && (axis > sm_ndim)))
{{
PyErr_SetString(PyExc_ValueError, "invalid axis in SoftmaxGrad");
{fail};
}}
}} }}
if (({dx} == NULL) if (({dx} == NULL)
|| !(PyArray_CompareLists(PyArray_DIMS({dx}), PyArray_DIMS({sm}), sm_ndim))) || !(PyArray_CompareLists(PyArray_DIMS({dx}), PyArray_DIMS({sm}), sm_ndim)))
{{ {{
...@@ -289,10 +296,14 @@ class Softmax(COp): ...@@ -289,10 +296,14 @@ class Softmax(COp):
def c_headers(self, **kwargs): def c_headers(self, **kwargs):
return ["<iostream>", "<cmath>"] return ["<iostream>", "<cmath>"]
def c_support_code_apply(self, node: Apply, name: str) -> str:
"""Needed to define NPY_RAVEL_AXIS"""
return npy_2_compat_header()
def c_code(self, node, name, inp, out, sub): def c_code(self, node, name, inp, out, sub):
(x,) = inp (x,) = inp
(sm,) = out (sm,) = out
axis = self.axis if self.axis is not None else np.MAXDIMS axis = self.axis if self.axis is not None else "NPY_RAVEL_AXIS"
fail = sub["fail"] fail = sub["fail"]
# dtype = node.inputs[0].type.dtype_specs()[1] # dtype = node.inputs[0].type.dtype_specs()[1]
# TODO: put this into a templated function, in the support code # TODO: put this into a templated function, in the support code
...@@ -309,7 +320,7 @@ class Softmax(COp): ...@@ -309,7 +320,7 @@ class Softmax(COp):
int x_ndim = PyArray_NDIM({x}); int x_ndim = PyArray_NDIM({x});
int axis = {axis}; int axis = {axis};
int iterate_axis = !(axis == NPY_MAXDIMS || x_ndim == 1); int iterate_axis = !(axis == NPY_RAVEL_AXIS || x_ndim == 1);
// Validate inputs // Validate inputs
if ((PyArray_TYPE({x}) != NPY_DOUBLE) && if ((PyArray_TYPE({x}) != NPY_DOUBLE) &&
...@@ -319,11 +330,14 @@ class Softmax(COp): ...@@ -319,11 +330,14 @@ class Softmax(COp):
{fail} {fail}
}} }}
if (axis < 0) axis = x_ndim + axis; if (iterate_axis)
if ((axis < 0) || (iterate_axis && (axis > x_ndim)))
{{ {{
PyErr_SetString(PyExc_ValueError, "invalid axis in Softmax"); if (axis < 0) axis = x_ndim + axis;
{fail} if ((axis < 0) || (iterate_axis && (axis > x_ndim)))
{{
PyErr_SetString(PyExc_ValueError, "invalid axis in Softmax");
{fail}
}}
}} }}
// Allocate Output Array // Allocate Output Array
...@@ -481,7 +495,7 @@ class Softmax(COp): ...@@ -481,7 +495,7 @@ class Softmax(COp):
@staticmethod @staticmethod
def c_code_cache_version(): def c_code_cache_version():
return (4,) return (5,)
def softmax(c, axis=None): def softmax(c, axis=None):
...@@ -541,10 +555,14 @@ class LogSoftmax(COp): ...@@ -541,10 +555,14 @@ class LogSoftmax(COp):
def c_headers(self, **kwargs): def c_headers(self, **kwargs):
return ["<cmath>"] return ["<cmath>"]
def c_support_code_apply(self, node: Apply, name: str) -> str:
"""Needed to define NPY_RAVEL_AXIS"""
return npy_2_compat_header()
def c_code(self, node, name, inp, out, sub): def c_code(self, node, name, inp, out, sub):
(x,) = inp (x,) = inp
(sm,) = out (sm,) = out
axis = self.axis if self.axis is not None else np.MAXDIMS axis = self.axis if self.axis is not None else "NPY_RAVEL_AXIS"
fail = sub["fail"] fail = sub["fail"]
return dedent( return dedent(
...@@ -558,7 +576,7 @@ class LogSoftmax(COp): ...@@ -558,7 +576,7 @@ class LogSoftmax(COp):
int x_ndim = PyArray_NDIM({x}); int x_ndim = PyArray_NDIM({x});
int axis = {axis}; int axis = {axis};
int iterate_axis = !(axis == NPY_MAXDIMS || x_ndim == 1); int iterate_axis = !(axis == NPY_RAVEL_AXIS || x_ndim == 1);
// Validate inputs // Validate inputs
if ((PyArray_TYPE({x}) != NPY_DOUBLE) && if ((PyArray_TYPE({x}) != NPY_DOUBLE) &&
...@@ -568,13 +586,15 @@ class LogSoftmax(COp): ...@@ -568,13 +586,15 @@ class LogSoftmax(COp):
{fail} {fail}
}} }}
if (axis < 0) axis = x_ndim + axis; if (iterate_axis)
if ((axis < 0) || (iterate_axis && (axis > x_ndim)))
{{ {{
PyErr_SetString(PyExc_ValueError, "invalid axis in LogSoftmax"); if (axis < 0) axis = x_ndim + axis;
{fail} if ((axis < 0) || (iterate_axis && (axis > x_ndim)))
{{
PyErr_SetString(PyExc_ValueError, "invalid axis in LogSoftmax");
{fail}
}}
}} }}
// Allocate Output Array // Allocate Output Array
if (({sm}) == NULL || !(PyArray_CompareLists(PyArray_DIMS({sm}), PyArray_DIMS({x}), x_ndim))) if (({sm}) == NULL || !(PyArray_CompareLists(PyArray_DIMS({sm}), PyArray_DIMS({x}), x_ndim)))
{{ {{
...@@ -730,7 +750,7 @@ class LogSoftmax(COp): ...@@ -730,7 +750,7 @@ class LogSoftmax(COp):
@staticmethod @staticmethod
def c_code_cache_version(): def c_code_cache_version():
return (1,) return (2,)
def log_softmax(c, axis=None): def log_softmax(c, axis=None):
......
...@@ -18,7 +18,7 @@ from pytensor.graph.type import Type ...@@ -18,7 +18,7 @@ from pytensor.graph.type import Type
from pytensor.graph.utils import MethodNotDefined from pytensor.graph.utils import MethodNotDefined
from pytensor.link.c.op import COp from pytensor.link.c.op import COp
from pytensor.link.c.params_type import ParamsType from pytensor.link.c.params_type import ParamsType
from pytensor.npy_2_compat import numpy_version, using_numpy_2 from pytensor.npy_2_compat import npy_2_compat_header, numpy_version, using_numpy_2
from pytensor.printing import Printer, pprint, set_precedence from pytensor.printing import Printer, pprint, set_precedence
from pytensor.scalar.basic import ScalarConstant, ScalarVariable from pytensor.scalar.basic import ScalarConstant, ScalarVariable
from pytensor.tensor import ( from pytensor.tensor import (
...@@ -2149,7 +2149,7 @@ class AdvancedSubtensor1(COp): ...@@ -2149,7 +2149,7 @@ class AdvancedSubtensor1(COp):
def c_support_code(self, **kwargs): def c_support_code(self, **kwargs):
# In some versions of numpy, NPY_MIN_INTP is defined as MIN_LONG, # In some versions of numpy, NPY_MIN_INTP is defined as MIN_LONG,
# which is not defined. It should be NPY_MIN_LONG instead in that case. # which is not defined. It should be NPY_MIN_LONG instead in that case.
return dedent( return npy_2_compat_header() + dedent(
"""\ """\
#ifndef MIN_LONG #ifndef MIN_LONG
#define MIN_LONG NPY_MIN_LONG #define MIN_LONG NPY_MIN_LONG
...@@ -2174,7 +2174,7 @@ class AdvancedSubtensor1(COp): ...@@ -2174,7 +2174,7 @@ class AdvancedSubtensor1(COp):
if (!PyArray_CanCastSafely(i_type, NPY_INTP) && if (!PyArray_CanCastSafely(i_type, NPY_INTP) &&
PyArray_SIZE({i_name}) > 0) {{ PyArray_SIZE({i_name}) > 0) {{
npy_int64 min_val, max_val; npy_int64 min_val, max_val;
PyObject* py_min_val = PyArray_Min({i_name}, NPY_MAXDIMS, PyObject* py_min_val = PyArray_Min({i_name}, NPY_RAVEL_AXIS,
NULL); NULL);
if (py_min_val == NULL) {{ if (py_min_val == NULL) {{
{fail}; {fail};
...@@ -2184,7 +2184,7 @@ class AdvancedSubtensor1(COp): ...@@ -2184,7 +2184,7 @@ class AdvancedSubtensor1(COp):
if (min_val == -1 && PyErr_Occurred()) {{ if (min_val == -1 && PyErr_Occurred()) {{
{fail}; {fail};
}} }}
PyObject* py_max_val = PyArray_Max({i_name}, NPY_MAXDIMS, PyObject* py_max_val = PyArray_Max({i_name}, NPY_RAVEL_AXIS,
NULL); NULL);
if (py_max_val == NULL) {{ if (py_max_val == NULL) {{
{fail}; {fail};
...@@ -2243,7 +2243,7 @@ class AdvancedSubtensor1(COp): ...@@ -2243,7 +2243,7 @@ class AdvancedSubtensor1(COp):
""" """
def c_code_cache_version(self): def c_code_cache_version(self):
return (0, 1, 2) return (0, 1, 2, 3)
advanced_subtensor1 = AdvancedSubtensor1() advanced_subtensor1 = AdvancedSubtensor1()
......
...@@ -18,6 +18,7 @@ from pytensor.graph.fg import FunctionGraph ...@@ -18,6 +18,7 @@ from pytensor.graph.fg import FunctionGraph
from pytensor.graph.replace import vectorize_node from pytensor.graph.replace import vectorize_node
from pytensor.link.basic import PerformLinker from pytensor.link.basic import PerformLinker
from pytensor.link.c.basic import CLinker, OpWiseCLinker from pytensor.link.c.basic import CLinker, OpWiseCLinker
from pytensor.npy_2_compat import numpy_maxdims
from pytensor.tensor import as_tensor_variable from pytensor.tensor import as_tensor_variable
from pytensor.tensor.basic import get_scalar_constant_value, second from pytensor.tensor.basic import get_scalar_constant_value, second
from pytensor.tensor.elemwise import CAReduce, DimShuffle, Elemwise from pytensor.tensor.elemwise import CAReduce, DimShuffle, Elemwise
...@@ -121,7 +122,8 @@ class TestDimShuffle(unittest_tools.InferShapeTester): ...@@ -121,7 +122,8 @@ class TestDimShuffle(unittest_tools.InferShapeTester):
def test_too_big_rank(self): def test_too_big_rank(self):
x = self.type(self.dtype, shape=())() x = self.type(self.dtype, shape=())()
y = x.dimshuffle(("x",) * (np.MAXDIMS + 1)) y = x.dimshuffle(("x",) * (numpy_maxdims + 1))
with pytest.raises(ValueError): with pytest.raises(ValueError):
y.eval({x: 0}) y.eval({x: 0})
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论