Skip to content

Numpy 2.x C implementation for AdvancedIncSubtensor1 #1325

Open
@ricardoV94

Description

@ricardoV94

Description

In our bump to numpy 2.0 we had to stop supporting the C implementation of this Op #1194

This was because we were using this PyArrayMapIter functionality that numpy removed. The Op needs to be able to iterate over the first dimension of an array to increment the values provided in y. This increment may need to broadcast with the selected entries.

Example:

import numpy as np
import pytensor
import pytensor.tensor as pt

x = pt.tensor("x", shape=(6, 2, 3))
y = pt.vector("y", shape=(3,))
idxs = pt.vector("idxs", shape=(7,), dtype=int)
out = pt.subtensor.advanced_inc_subtensor1(x, y, idxs)

x_test = np.zeros(x.type.shape)
y_test = np.ones(y.type.shape)
idxs_test = [0, 0, 0, 1, 2, 3, 5]
res = out.eval({x: x_test, y: y_test, idxs: idxs_test})
print(res)
# [[[3. 3. 3.]
#   [3. 3. 3.]]
#  [[1. 1. 1.]
#   [1. 1. 1.]]
#  [[1. 1. 1.]
#   [1. 1. 1.]]
#  [[1. 1. 1.]
#   [1. 1. 1.]]
#  [[0. 0. 0.]
#   [0. 0. 0.]]
#  [[1. 1. 1.]
#   [1. 1. 1.]]]

It's a restricted version of numpy.add.at, where indices must be a single integer vector, whereas numpy.add.at supports arbitrary advanced indexing.

res_np = x_test.copy()
np.add.at(res_np, idxs_test, y_test)
np.testing.assert_allclose(res, res_np)

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions