mirror of
https://github.com/ROCm/jax.git
synced 2025-04-14 10:56:06 +00:00
293 lines
9.2 KiB
Python
293 lines
9.2 KiB
Python
# Copyright 2020 The JAX Authors.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# https://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from functools import partial
|
|
|
|
from absl.testing import absltest
|
|
import numpy as np
|
|
|
|
import jax
|
|
from jax import numpy as jnp
|
|
from jax._src import test_util as jtu
|
|
|
|
jax.config.parse_flags_with_absl()
|
|
|
|
|
|
class VectorizeTest(jtu.JaxTestCase):
|
|
|
|
@jtu.sample_product(
|
|
[dict(left_shape=left_shape, right_shape=right_shape,
|
|
result_shape=result_shape)
|
|
for left_shape, right_shape, result_shape in [
|
|
((2, 3), (3, 4), (2, 4)),
|
|
((2, 3), (1, 3, 4), (1, 2, 4)),
|
|
((1, 2, 3), (1, 3, 4), (1, 2, 4)),
|
|
((5, 2, 3), (1, 3, 4), (5, 2, 4)),
|
|
((6, 5, 2, 3), (3, 4), (6, 5, 2, 4)),
|
|
]
|
|
],
|
|
)
|
|
@jax.numpy_rank_promotion('allow')
|
|
def test_matmat(self, left_shape, right_shape, result_shape):
|
|
matmat = jnp.vectorize(jnp.dot, signature='(n,m),(m,k)->(n,k)')
|
|
self.assertEqual(matmat(jnp.zeros(left_shape),
|
|
jnp.zeros(right_shape)).shape, result_shape)
|
|
|
|
@jtu.sample_product(
|
|
[dict(left_shape=left_shape, right_shape=right_shape,
|
|
result_shape=result_shape)
|
|
for left_shape, right_shape, result_shape in [
|
|
((2, 3), (3,), (2,)),
|
|
((2, 3), (1, 3), (1, 2)),
|
|
((4, 2, 3), (1, 3), (4, 2)),
|
|
((5, 4, 2, 3), (1, 3), (5, 4, 2)),
|
|
]
|
|
],
|
|
)
|
|
@jax.numpy_rank_promotion('allow')
|
|
def test_matvec(self, left_shape, right_shape, result_shape):
|
|
matvec = jnp.vectorize(jnp.dot, signature='(n,m),(m)->(n)')
|
|
self.assertEqual(matvec(jnp.zeros(left_shape),
|
|
jnp.zeros(right_shape)).shape, result_shape)
|
|
|
|
@jtu.sample_product(
|
|
[dict(left_shape=left_shape, right_shape=right_shape,
|
|
result_shape=result_shape)
|
|
for left_shape, right_shape, result_shape in [
|
|
((3,), (3,), ()),
|
|
((2, 3), (3,), (2,)),
|
|
((4, 2, 3), (3,), (4, 2)),
|
|
]
|
|
],
|
|
)
|
|
@jax.numpy_rank_promotion('allow')
|
|
def test_vecmat(self, left_shape, right_shape, result_shape):
|
|
vecvec = jnp.vectorize(jnp.dot, signature='(m),(m)->()')
|
|
self.assertEqual(vecvec(jnp.zeros(left_shape),
|
|
jnp.zeros(right_shape)).shape, result_shape)
|
|
|
|
@jtu.sample_product(
|
|
[dict(shape=shape, result_shape=result_shape)
|
|
for shape, result_shape in [
|
|
((3,), ()),
|
|
((2, 3,), (2,)),
|
|
((1, 2, 3,), (1, 2)),
|
|
]
|
|
],
|
|
)
|
|
def test_magnitude(self, shape, result_shape):
|
|
size = 1
|
|
for x in shape:
|
|
size *= x
|
|
inputs = jnp.arange(size).reshape(shape)
|
|
|
|
@partial(jnp.vectorize, signature='(n)->()')
|
|
def magnitude(x):
|
|
return jnp.dot(x, x)
|
|
|
|
self.assertEqual(magnitude(inputs).shape, result_shape)
|
|
|
|
@jtu.sample_product(
|
|
[dict(shape=shape, result_shape=result_shape)
|
|
for shape, result_shape in [
|
|
((3,), ()),
|
|
((2, 3), (2,)),
|
|
((1, 2, 3, 4), (1, 2, 3)),
|
|
]
|
|
],
|
|
)
|
|
def test_mean(self, shape, result_shape):
|
|
mean = jnp.vectorize(jnp.mean, signature='(n)->()')
|
|
self.assertEqual(mean(jnp.zeros(shape)).shape, result_shape)
|
|
|
|
@jtu.sample_product(
|
|
[dict(shape=shape, result_shape=result_shape)
|
|
for shape, result_shape in [
|
|
((), (2,)),
|
|
((3,), (3,2,)),
|
|
]
|
|
],
|
|
)
|
|
def test_stack_plus_minus(self, shape, result_shape):
|
|
|
|
@partial(jnp.vectorize, signature='()->(n)')
|
|
def stack_plus_minus(x):
|
|
return jnp.stack([x, -x])
|
|
|
|
self.assertEqual(stack_plus_minus(jnp.zeros(shape)).shape, result_shape)
|
|
|
|
def test_center(self):
|
|
|
|
@partial(jnp.vectorize, signature='(n)->(),(n)')
|
|
def center(array):
|
|
bias = jnp.mean(array)
|
|
debiased = array - bias
|
|
return bias, debiased
|
|
|
|
b, a = center(jnp.arange(3.0))
|
|
self.assertEqual(a.shape, (3,))
|
|
self.assertEqual(b.shape, ())
|
|
self.assertAllClose(1.0, b, check_dtypes=False)
|
|
|
|
b, a = center(jnp.arange(6.0).reshape(2, 3))
|
|
self.assertEqual(a.shape, (2, 3))
|
|
self.assertEqual(b.shape, (2,))
|
|
self.assertAllClose(jnp.array([1.0, 4.0]), b, check_dtypes=False)
|
|
|
|
def test_exclude_first(self):
|
|
|
|
@partial(jnp.vectorize, excluded={0})
|
|
def f(x, y):
|
|
assert x == 'foo'
|
|
assert y.ndim == 0
|
|
return y
|
|
|
|
x = jnp.arange(3)
|
|
self.assertAllClose(x, f('foo', x))
|
|
self.assertAllClose(x, jax.jit(f, static_argnums=0)('foo', x))
|
|
|
|
def test_exclude_second(self):
|
|
|
|
@partial(jnp.vectorize, excluded={1})
|
|
def f(x, y):
|
|
assert x.ndim == 0
|
|
assert y == 'foo'
|
|
return x
|
|
|
|
x = jnp.arange(3)
|
|
self.assertAllClose(x, f(x, 'foo'))
|
|
self.assertAllClose(x, jax.jit(f, static_argnums=1)(x, 'foo'))
|
|
|
|
def test_exclude_kwargs(self):
|
|
@partial(np.vectorize, excluded=(2, 'func'))
|
|
def f_np(x, y, func=np.add):
|
|
assert np.ndim(x) == np.ndim(y) == 0
|
|
return func(x, y)
|
|
|
|
@partial(jnp.vectorize, excluded=(2, 'func'))
|
|
def f_jnp(x, y, func=jnp.add):
|
|
assert x.ndim == y.ndim == 0
|
|
return func(x, y)
|
|
|
|
x = np.arange(4, dtype='int32')
|
|
y = np.int32(2)
|
|
|
|
self.assertArraysEqual(f_np(x, y), f_jnp(x, y))
|
|
self.assertArraysEqual(f_np(x, y, np.power), f_jnp(x, y, jnp.power))
|
|
self.assertArraysEqual(f_np(x, y, func=np.power), f_jnp(x, y, func=jnp.power))
|
|
|
|
def test_exclude_errors(self):
|
|
with self.assertRaisesRegex(
|
|
TypeError, "jax.numpy.vectorize can only exclude"):
|
|
jnp.vectorize(lambda x: x, excluded={1.5})
|
|
|
|
with self.assertRaisesRegex(
|
|
ValueError, r"excluded=\{-1\} contains negative numbers"):
|
|
jnp.vectorize(lambda x: x, excluded={-1})
|
|
|
|
def test_bad_inputs(self):
|
|
matmat = jnp.vectorize(jnp.dot, signature='(n,m),(m,k)->(n,k)')
|
|
with self.assertRaisesRegex(
|
|
TypeError, "wrong number of positional arguments"):
|
|
matmat(jnp.zeros((3, 2)))
|
|
with self.assertRaisesRegex(
|
|
ValueError,
|
|
r"input with shape \(2,\) does not have enough dimensions"):
|
|
matmat(jnp.zeros((2,)), jnp.zeros((2, 2)))
|
|
with self.assertRaisesRegex(
|
|
ValueError, r"inconsistent size for core dimension 'm'"):
|
|
matmat(jnp.zeros((2, 3)), jnp.zeros((4, 5)))
|
|
|
|
def test_wrong_output_type(self):
|
|
f = jnp.vectorize(jnp.dot, signature='(n,m),(m,k)->(n,k),()')
|
|
with self.assertRaisesRegex(
|
|
TypeError, "output must be a tuple"):
|
|
f(jnp.zeros((2, 2)), jnp.zeros((2, 2)))
|
|
|
|
def test_wrong_num_outputs(self):
|
|
f = jnp.vectorize(lambda *args: args, signature='(),()->(),(),()')
|
|
with self.assertRaisesRegex(
|
|
TypeError, "wrong number of output arguments"):
|
|
f(1, 2)
|
|
|
|
def test_wrong_output_shape(self):
|
|
f = jnp.vectorize(jnp.dot, signature='(n,m),(m,k)->(n)')
|
|
with self.assertRaisesRegex(
|
|
ValueError, r"output shape \(2, 2\) does not match"):
|
|
f(jnp.zeros((2, 2)), jnp.zeros((2, 2)))
|
|
|
|
def test_inconsistent_output_size(self):
|
|
f = jnp.vectorize(jnp.dot, signature='(n,m),(m,k)->(n,n)')
|
|
with self.assertRaisesRegex(
|
|
ValueError, r"inconsistent size for core dimension 'n'"):
|
|
f(jnp.zeros((2, 3)), jnp.zeros((3, 4)))
|
|
|
|
def test_expand_dims_multiple_outputs_no_signature(self):
|
|
f = jnp.vectorize(lambda x: (x, x))
|
|
x = jnp.arange(1)
|
|
xx = f(x)
|
|
self.assertAllClose(xx[0], x)
|
|
self.assertAllClose(xx[1], x)
|
|
self.assertIsInstance(xx, tuple)
|
|
|
|
def test_none_arg(self):
|
|
f = jnp.vectorize(lambda x, y: x if y is None else x + y)
|
|
x = jnp.arange(10)
|
|
self.assertAllClose(f(x, None), x)
|
|
|
|
y = jnp.arange(10, 20)
|
|
self.assertAllClose(f(x, y), x + y)
|
|
|
|
def test_none_arg_bad_signature(self):
|
|
f = jnp.vectorize(lambda x, y: x if y is None else x + y,
|
|
signature='(k),(k)->(k)')
|
|
args = jnp.arange(10), None
|
|
msg = r"Cannot pass None at locations \{1\} with signature='\(k\),\(k\)->\(k\)'"
|
|
with self.assertRaisesRegex(ValueError, msg):
|
|
f(*args)
|
|
|
|
def test_rank_promotion_error(self):
|
|
# Regression test for https://github.com/jax-ml/jax/issues/22305
|
|
f = jnp.vectorize(jnp.add, signature="(),()->()")
|
|
rank2 = jnp.zeros((10, 10))
|
|
rank1 = jnp.zeros(10)
|
|
rank0 = jnp.zeros(())
|
|
msg = "operands with shapes .* require rank promotion"
|
|
with jax.numpy_rank_promotion('raise'):
|
|
with self.assertRaisesRegex(ValueError, msg):
|
|
f(rank2, rank1)
|
|
with jax.numpy_rank_promotion('warn'):
|
|
with self.assertWarnsRegex(UserWarning, msg):
|
|
f(rank2, rank1)
|
|
|
|
# no warning for scalar rank promotion
|
|
with jax.numpy_rank_promotion('raise'):
|
|
f(rank2, rank0)
|
|
f(rank1, rank0)
|
|
with jax.numpy_rank_promotion('warn'):
|
|
f(rank2, rank0)
|
|
f(rank1, rank0)
|
|
|
|
# No warning when broadcasted ranks match.
|
|
f2 = jnp.vectorize(jnp.add, signature="(n),()->(n)")
|
|
with jax.numpy_rank_promotion('raise'):
|
|
f2(rank2, rank1)
|
|
with jax.numpy_rank_promotion('warn'):
|
|
with self.assertNoWarnings():
|
|
f2(rank2, rank1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
absltest.main(testLoader=jtu.JaxTestLoader())
|