Buffer Changes

Author: Dhruvanshu-Joshi

pytensor/graph/op.py

@@ -291,6 +291,7 @@ def __call__(
 
         """
         node = self.make_node(*inputs, **kwargs)
+
         if name is not None:
             if len(node.outputs) == 1:
                 node.outputs[0].name = name

pytensor/ifelse.py

@@ -477,7 +477,7 @@ def cond_make_inplace(fgraph, node):
     Reshape,
     Unbroadcast,
     pt.math.Dot,
-    pt.math.TensorMax,
+    pt.math.Max,
     pt.math.Argmax,
     pt.subtensor.Subtensor,
     pt.subtensor.IncSubtensor,

pytensor/tensor/math.py

@@ -808,7 +808,8 @@ def max_and_argmax(a, axis=None, keepdims=False):
     axis = check_and_normalize_axes(a, axis)
     if len(axis) == 0:
         axis = list(range(a.type.ndim))
-    out = TensorMax(axis)(a)
+    # out = TensorMax(axis)(a)
+    out = Max(axis)(a)
     argout = Argmax(axis)(a)
     # out, argout = MaxAndArgmax(axis)(a)
 
@@ -864,6 +865,105 @@ def clone(self, **kwargs):
         axis = kwargs.get("axis", self.axis)
         return type(self)(axis=axis)
 
+    def grad(self, inp, grads):
+        # The strict sense mathematical gradient of the maximum function is
+        # not calculated here for it is not defined at every point where some
+        # coordinates are identical. However, since the latter set has null
+        # Lebesgue measure, the result may be interpreted as weak gradient.
+
+        # @note: This function should work correctly for L{vector}s.
+        # (x, y), (gz, gw)
+        # gz*dz/dx + gw*dw/dx, gz*dz/dy + gw*dw/dy
+        # gMax * dMax/dx + gArgMax * dArgMax/dx,
+        # gMax * dMax/daxis + gArgMax * dArgMax/daxis
+        # g_max has one less dimension than x, so you need to complete
+        # g_max to x's shape when axis=0 the broadcasting mechanism
+        # does it automatically
+        x = inp[0]
+        axis = as_tensor_variable(self.axis)
+        (g_max,) = grads
+
+        g_max_disconnected = isinstance(g_max.type, DisconnectedType)
+
+        # if the op is totally disconnected, so are its inputs
+        if g_max_disconnected:
+            return [DisconnectedType()()]
+
+        if NoneConst.equals(axis):
+            axis_ = list(range(x.ndim))
+        else:
+            axis_ = axis
+        xmax = max(x, axis_)
+
+        # Raise the g_max and xmax to the same number of dim as the input.
+        pattern = []
+        out_dim = 0
+        if NoneConst.equals(axis):
+            # We are taking the max/argmax over all dimensions.
+            axis = None
+        for i in range(x.ndim):
+            if axis is None or i in axis.data:
+                pattern.append("x")
+            else:
+                pattern.append(out_dim)
+                out_dim += 1
+        g_max_pad = DimShuffle(g_max.broadcastable, pattern)(g_max)
+        xmax_pad = DimShuffle(xmax.broadcastable, pattern)(xmax)
+
+        # Set the grad to the correct position.
+        g_x = eq(xmax_pad, x) * g_max_pad
+        return (g_x,)
+
+    def c_code(self, node, name, inp, out, sub):
+        if len(self.axis) != 1 and len(self.axis) != node.inputs[0].ndim:
+            raise NotImplementedError(
+                "NumPy C-API can compute max only for 1 axis or for all axes."
+            )
+        x = inp[0]
+        axis = sub["params"]
+        # max, argmax = out
+        (max,) = out
+        fail = sub["fail"]
+        ret = """
+        #if PY_MAJOR_VERSION >= 3
+            #ifndef PyInt_AS_LONG
+                #define PyInt_AS_LONG PyLong_AS_LONG
+            #endif
+        #endif
+
+        int axis;
+
+        if (PyTuple_GET_SIZE(%(axis)s) == PyArray_NDIM(%(x)s)) {
+            axis = NPY_MAXDIMS;
+        } else if(PyTuple_GET_SIZE(%(axis)s) == 1) {
+            PyObject* axis_object = PyTuple_GET_ITEM(%(axis)s, 0);
+            axis = (int)PyInt_AS_LONG(axis_object);
+            if (axis > PyArray_NDIM(%(x)s)-1 || axis < -PyArray_NDIM(%(x)s)) {
+                PyErr_SetString(PyExc_ValueError,
+                "TensorMax: bad axis argument");
+                %(fail)s
+            }
+        } else {
+            PyErr_SetString(PyExc_NotImplementedError,
+            "TensorMax: NumPy C-API can compute max only for 1 axis or for all axes.");
+            %(fail)s
+        }
+
+        Py_CLEAR(%(max)s);
+
+        %(max)s = (PyArrayObject*)PyArray_Max(%(x)s, axis, NULL);
+        if (%(max)s == NULL) {
+            %(fail)s;
+        }
+        if (!PyArray_CheckExact(%(max)s)) {
+            %(max)s = (PyArrayObject*)PyArray_FromAny((PyObject*)%(max)s, NULL, 0, 0, NPY_ARRAY_ENSUREARRAY, NULL);
+            if(%(max)s == NULL){
+                %(fail)s;
+            }
+        }
+        """
+        return ret % locals()
+
 
 class Min(NonZeroDimsCAReduce):
     nfunc_spec = ("min", 1, 1)
@@ -956,6 +1056,12 @@ def min(x, axis=None, keepdims=False):
     elif str_x_type in uint_dtypes:
         itype = np.iinfo(x.dtype)
         max_val = np.array(itype.max, dtype=itype.dtype)
+        # print('a')
+        # for c in (max_val - x):
+        #     print(c.eval())
+        # print()
+        # print(max(max_val - x, axis=axis, keepdims=keepdims).eval())
+        # print()
         return max_val - max(max_val - x, axis=axis, keepdims=keepdims)
     elif str_x_type == "bool":
         return ~max(~x, axis=axis, keepdims=keepdims)

pytensor/tensor/rewriting/uncanonicalize.py

@@ -31,10 +31,11 @@
 
 """
 
+from pytensor import scalar as ps
 from pytensor.graph.rewriting.basic import copy_stack_trace, node_rewriter
 from pytensor.tensor.basic import Alloc, alloc, constant
-from pytensor.tensor.elemwise import DimShuffle
-from pytensor.tensor.math import Min, TensorMax, neg
+from pytensor.tensor.elemwise import CAReduce, DimShuffle
+from pytensor.tensor.math import Min, neg
 from pytensor.tensor.rewriting.basic import register_uncanonicalize
 from pytensor.tensor.shape import Reshape, reshape
 from pytensor.tensor.subtensor import Subtensor
@@ -54,13 +55,14 @@ def local_max_to_min(fgraph, node):
     the interface put only MaxAndArgmax into the graph.
 
     """
-    # pytensor.dprint(node)
-    # print()
-    # print(node.op == neg)
     if node.op == neg and node.inputs[0].owner:
         max = node.inputs[0]
         # print(max.owner.op.scalar_op)
-        if max.owner and isinstance(max.owner.op, TensorMax):
+        if (
+            max.owner
+            and isinstance(max.owner.op, CAReduce)
+            and max.owner.op.scalar_op == ps.scalar_maximum
+        ):
             neg_node = max.owner.inputs[0]
             if neg_node.owner and neg_node.owner.op == neg:
                 new = Min(max.owner.op.axis)(neg_node.owner.inputs[0])

tests/tensor/rewriting/test_uncanonicalize.py

@@ -9,7 +9,6 @@
 from pytensor.graph.rewriting.basic import out2in
 from pytensor.link.basic import PerformLinker
 from pytensor.tensor.elemwise import CAReduce, DimShuffle, Elemwise
-from pytensor.tensor.math import TensorMax
 from pytensor.tensor.math import min as pt_min
 from pytensor.tensor.rewriting.uncanonicalize import (
     local_alloc_dimshuffle,
@@ -43,7 +42,7 @@ def test_optimization_min(self):
             f = function([n], pt_min(-n, axis), mode=self.mode)
             topo = f.maker.fgraph.toposort()
             assert len(topo) == 2
-            assert isinstance(topo[0].op, TensorMax)  # max
+            assert isinstance(topo[0].op, CAReduce)  # max
             assert isinstance(topo[1].op, Elemwise)
             assert isinstance(topo[1].op.scalar_op, ps.Neg)
             f(data)
@@ -53,13 +52,13 @@ def test_optimization_min(self):
             assert len(topo) == 2
             assert isinstance(topo[0].op, Elemwise)
             assert isinstance(topo[0].op.scalar_op, ps.Neg)
-            assert isinstance(topo[1].op, TensorMax)  # max
+            assert isinstance(topo[1].op, CAReduce)  # max
             f(data)
 
             f = function([n], -pt_min(-n, axis), mode=self.mode)
             topo = f.maker.fgraph.toposort()
             assert len(topo) == 1
-            assert isinstance(topo[0].op, TensorMax)  # max
+            assert isinstance(topo[0].op, CAReduce)  # max
             f(data)
 
 

tests/tensor/test_max_argmax.py

@@ -16,7 +16,7 @@
 )
 from pytensor.tensor.math import (
     Argmax,
-    TensorMax,
+    Max,
     argmax,
     argmin,
     max,
@@ -43,7 +43,7 @@
 
 class TestMaxAndArgmax:
     def setup_method(self):
-        TensorMax.debug = 0
+        Max.debug = 0
 
     def test_basic(self):
         # dbt: for some reason, Argmax does not work when I pass: n = as_tensor_variable(5.0)
@@ -317,7 +317,8 @@ def test_vectorize(self, core_axis, batch_axis):
         # Test MaxAndArgmax
         max_x, argmax_x = max_and_argmax(x, axis=core_axis)
         node = max_x.owner
-        assert isinstance(node.op, TensorMax)
+        # assert isinstance(node.op, TensorMax)
+        assert isinstance(node.op, Max)
 
         # dbt: how to make Argmax user facing?
         # new_node = vectorize_node(node, batch_x)
@@ -339,7 +340,7 @@ def test_vectorize(self, core_axis, batch_axis):
 
 class TestArgminArgmax:
     def setup_method(self):
-        TensorMax.debug = 0
+        Argmax.debug = 0
 
     def test_scalar(self):
         for fct in [argmin, argmax]:
@@ -501,7 +502,7 @@ def test_bool(self):
 
 class TestMinMax:
     def setup_method(self):
-        TensorMax.debug = 0
+        Max.debug = 0
 
     def test_scalar(self):
         for fct in [max, min]:
@@ -675,6 +676,12 @@ def test_uint(self):
             n = as_tensor_variable(data)
             assert min(n).dtype == dtype
             i = eval_outputs(min(n))
+            # pytensor.dprint(n)
+            for x in n:
+                print(x.eval())
+            print(i)
+            print(itype.min)
+            print()
             assert i == itype.min
             assert max(n).dtype == dtype
             i = eval_outputs(max(n))