Skip to content
63 changes: 2 additions & 61 deletions dpnp/dpnp_iface_linearalgebra.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,9 +37,7 @@

"""


import numpy
from dpctl.tensor._numpy_helper import normalize_axis_tuple

import dpnp

Expand All @@ -48,6 +46,7 @@
dpnp_dot,
dpnp_kron,
dpnp_matmul,
dpnp_tensordot,
dpnp_vecdot,
)

Expand Down Expand Up @@ -1047,65 +1046,7 @@ def tensordot(a, b, axes=2):
# TODO: use specific scalar-vector kernel
return dpnp.multiply(a, b)

try:
iter(axes)
except Exception as e: # pylint: disable=broad-exception-caught
if not isinstance(axes, int):
raise TypeError("Axes must be an integer.") from e
if axes < 0:
raise ValueError("Axes must be a non-negative integer.") from e
axes_a = tuple(range(-axes, 0))
axes_b = tuple(range(0, axes))
else:
if len(axes) != 2:
raise ValueError("Axes must consist of two sequences.")

axes_a, axes_b = axes
axes_a = (axes_a,) if dpnp.isscalar(axes_a) else axes_a
axes_b = (axes_b,) if dpnp.isscalar(axes_b) else axes_b

if len(axes_a) != len(axes_b):
raise ValueError("Axes length mismatch.")

# Make the axes non-negative
a_ndim = a.ndim
b_ndim = b.ndim
axes_a = normalize_axis_tuple(axes_a, a_ndim, "axis_a")
axes_b = normalize_axis_tuple(axes_b, b_ndim, "axis_b")

if a.ndim == 0 or b.ndim == 0:
# TODO: use specific scalar-vector kernel
return dpnp.multiply(a, b)

a_shape = a.shape
b_shape = b.shape
for axis_a, axis_b in zip(axes_a, axes_b):
if a_shape[axis_a] != b_shape[axis_b]:
raise ValueError(
"shape of input arrays is not similar at requested axes."
)

# Move the axes to sum over, to the end of "a"
not_in = tuple(k for k in range(a_ndim) if k not in axes_a)
newaxes_a = not_in + axes_a
n1 = int(numpy.prod([a_shape[ax] for ax in not_in]))
n2 = int(numpy.prod([a_shape[ax] for ax in axes_a]))
newshape_a = (n1, n2)
olda = [a_shape[axis] for axis in not_in]

# Move the axes to sum over, to the front of "b"
not_in = tuple(k for k in range(b_ndim) if k not in axes_b)
newaxes_b = tuple(axes_b + not_in)
n1 = int(numpy.prod([b_shape[ax] for ax in axes_b]))
n2 = int(numpy.prod([b_shape[ax] for ax in not_in]))
newshape_b = (n1, n2)
oldb = [b_shape[axis] for axis in not_in]

at = dpnp.transpose(a, newaxes_a).reshape(newshape_a)
bt = dpnp.transpose(b, newaxes_b).reshape(newshape_b)
res = dpnp.matmul(at, bt)

return res.reshape(olda + oldb)
return dpnp_tensordot(a, b, axes=axes)


def vdot(a, b):
Expand Down
73 changes: 72 additions & 1 deletion dpnp/dpnp_utils/dpnp_utils_linearalgebra.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,14 @@
from dpnp.dpnp_array import dpnp_array
from dpnp.dpnp_utils import get_usm_allocations

__all__ = ["dpnp_cross", "dpnp_dot", "dpnp_kron", "dpnp_matmul", "dpnp_vecdot"]
__all__ = [
"dpnp_cross",
"dpnp_dot",
"dpnp_kron",
"dpnp_matmul",
"dpnp_tensordot",
"dpnp_vecdot",
]


def _compute_res_dtype(*arrays, sycl_queue, dtype=None, casting="no"):
Expand Down Expand Up @@ -974,6 +981,70 @@ def dpnp_matmul(
return result


def dpnp_tensordot(a, b, axes=2):
"""Tensor dot product of two arrays."""

try:
iter(axes)
except Exception as e: # pylint: disable=broad-exception-caught
if not isinstance(axes, int):
raise TypeError("Axes must be an integer.") from e
if axes < 0:
raise ValueError("Axes must be a non-negative integer.") from e
axes_a = tuple(range(-axes, 0))
axes_b = tuple(range(0, axes))
else:
if len(axes) != 2:
raise ValueError("Axes must consist of two sequences.")

axes_a, axes_b = axes
axes_a = (axes_a,) if dpnp.isscalar(axes_a) else axes_a
axes_b = (axes_b,) if dpnp.isscalar(axes_b) else axes_b

if len(axes_a) != len(axes_b):
raise ValueError("Axes length mismatch.")

# Make the axes non-negative
a_ndim = a.ndim
b_ndim = b.ndim
axes_a = normalize_axis_tuple(axes_a, a_ndim, "axis_a")
axes_b = normalize_axis_tuple(axes_b, b_ndim, "axis_b")

if a.ndim == 0 or b.ndim == 0:
# TODO: use specific scalar-vector kernel
return dpnp.multiply(a, b)

a_shape = a.shape
b_shape = b.shape
for axis_a, axis_b in zip(axes_a, axes_b):
if a_shape[axis_a] != b_shape[axis_b]:
raise ValueError(
"shape of input arrays is not similar at requested axes."
)

# Move the axes to sum over, to the end of "a"
not_in = tuple(k for k in range(a_ndim) if k not in axes_a)
newaxes_a = not_in + axes_a
n1 = int(numpy.prod([a_shape[ax] for ax in not_in]))
n2 = int(numpy.prod([a_shape[ax] for ax in axes_a]))
newshape_a = (n1, n2)
olda = [a_shape[axis] for axis in not_in]

# Move the axes to sum over, to the front of "b"
not_in = tuple(k for k in range(b_ndim) if k not in axes_b)
newaxes_b = tuple(axes_b + not_in)
n1 = int(numpy.prod([b_shape[ax] for ax in axes_b]))
n2 = int(numpy.prod([b_shape[ax] for ax in not_in]))
newshape_b = (n1, n2)
oldb = [b_shape[axis] for axis in not_in]

at = dpnp.transpose(a, newaxes_a).reshape(newshape_a)
bt = dpnp.transpose(b, newaxes_b).reshape(newshape_b)
res = dpnp.matmul(at, bt)

return res.reshape(olda + oldb)


def dpnp_vecdot(
x1,
x2,
Expand Down
24 changes: 16 additions & 8 deletions dpnp/tests/helper.py
Original file line number Diff line number Diff line change
Expand Up @@ -162,7 +162,7 @@ def get_all_dtypes(


def generate_random_numpy_array(
shape, dtype=None, hermitian=False, seed_value=None
shape, dtype=None, hermitian=False, seed_value=None, low=-10, high=10
):
"""
Generate a random numpy array with the specified shape and dtype.
Expand All @@ -177,13 +177,19 @@ def generate_random_numpy_array(
dtype : str or dtype, optional
Desired data-type for the output array.
If not specified, data type will be determined by numpy.
Default : None
Default : ``None``
hermitian : bool, optional
If True, generates a Hermitian (symmetric if `dtype` is real) matrix.
Default : False
Default : ``False``
seed_value : int, optional
The seed value to initialize the random number generator.
Default : None
Default : ``None``
low : {int, float}, optional
Lower boundary of the generated samples from a uniform distribution.
Default : ``-10``.
high : {int, float}, optional
Upper boundary of the generated samples from a uniform distribution.
Default : ``10``.

Returns
-------
Expand All @@ -197,13 +203,15 @@ def generate_random_numpy_array(

"""

numpy.random.seed(seed_value)
numpy.random.seed(seed_value) if seed_value else numpy.random.seed(42)

a = numpy.random.randn(*shape).astype(dtype)
# dtype=int is needed for 0d arrays
size = numpy.prod(shape, dtype=int)
a = numpy.random.uniform(low, high, size).astype(dtype)
if numpy.issubdtype(a.dtype, numpy.complexfloating):
numpy.random.seed(seed_value)
a += 1j * numpy.random.randn(*shape)
a += 1j * numpy.random.uniform(low, high, size)

a = a.reshape(shape)
if hermitian and a.size > 0:
if a.ndim > 2:
orig_shape = a.shape
Expand Down
Loading
Loading