def testShapeAttribute_HasCorrectLength(self):
with utils.SaveCodeAsString() as code_saver:
x0 = jnp.zeros(())
x1 = jnp.zeros((1, ))
x2 = jnp.zeros((1, 2))
x3 = jnp.zeros((1, 2, 3))
x4 = jnp.zeros((1, 2, 3, 4))
x0_shape = x0.shape # pylint: disable=unused-variable
x1_shape = x1.shape # pylint: disable=unused-variable
x2_shape = x2.shape # pylint: disable=unused-variable
x3_shape = x3.shape # pylint: disable=unused-variable
x4_shape = x4.shape # pylint: disable=unused-variable
inferred = utils.pytype_infer_types(_PREAMBLE + code_saver.code)
self.assertEqual(x0_shape, ())
self.assertEqual(x1_shape, (1, ))
self.assertEqual(x2_shape, (1, 2))
self.assertEqual(x3_shape, (1, 2, 3))
self.assertEqual(x4_shape, (1, 2, 3, 4))
self.assertEqual('Tuple[()]', inferred.x0_shape)
self.assertEqual('Tuple[int]', inferred.x1_shape)
self.assertEqual('Tuple[int, int]', inferred.x2_shape)
self.assertEqual('Tuple[int, int, int]', inferred.x3_shape)
self.assertEqual('Tuple[int, int, int, int]', inferred.x4_shape)
def testShapeAttribute_HasLen(self):
with utils.SaveCodeAsString() as code_saver:
x = tf.zeros((1, ))
rank = len(x.shape) # pylint: disable=unused-variable
inferred = utils.pytype_infer_types(_PREAMBLE + code_saver.code)
self.assertEqual('int', inferred.rank)
def testZerosOnes_ReturnsCustomType(self):
with utils.SaveCodeAsString() as code_saver:
a = tf.zeros((1, )) # pylint: disable=unused-variable
b = tf.ones((1, )) # pylint: disable=unused-variable
inferred = utils.pytype_infer_types(_PREAMBLE + code_saver.code)
self.assertEqual(inferred.a, 'Tensor1')
self.assertEqual(inferred.b, 'Tensor1')
def testTensorAdd_ReturnsCustomType(self):
with utils.SaveCodeAsString() as code_saver:
x: Tensor1[A1] = tf.zeros((1, ))
a = x + 1 # pylint: disable=unused-variable
b = x + x # pylint: disable=unused-variable
inferred = utils.pytype_infer_types(_PREAMBLE + code_saver.code)
self.assertEqual(inferred.a, 'Tensor1[A1]')
self.assertEqual(inferred.b, 'Tensor1[A1]')
def testSumKeepdimsTrue_ReturnsAny(self):
# We haven't got around to making stubs for keepdims=True yet;
# make sure the type reflects that.
with utils.SaveCodeAsString() as code_saver:
x: Array1[A1] = jnp.zeros((1, ))
a = jnp.sum(x, axis=0, keepdims=True) # pylint: disable=unused-variable
b = jnp.sum(x, keepdims=True) # pylint: disable=unused-variable
inferred = utils.pytype_infer_types(_PREAMBLE + code_saver.code)
self.assertEqual(inferred.a, 'Any')
self.assertEqual(inferred.b, 'Any')
def testTensorUnaryOp_ReturnsCorrectTypeAndShape(self):
with utils.SaveCodeAsString() as code_saver:
x1: Tensor0 = tf.zeros(())
y1 = abs(x1) # pylint: disable=unused-variable
y2 = -x1 # pylint: disable=unused-variable
x2: Tensor1[A1] = tf.zeros((1, ))
y3 = abs(x2) # pylint: disable=unused-variable
y4 = -x2 # pylint: disable=unused-variable
inferred = utils.pytype_infer_types(_PREAMBLE + code_saver.code)
self.assertEqual('Tensor0', inferred.y1)
self.assertEqual('Tensor0', inferred.y2)
self.assertEqual('Tensor1[A1]', inferred.y3)
self.assertEqual('Tensor1[A1]', inferred.y4)
def testZerosOnes_ReturnsCorrectShape(self):
with utils.SaveCodeAsString() as code_saver:
a = jnp.zeros(()) # pylint: disable=unused-variable
b = jnp.ones(()) # pylint: disable=unused-variable
c = jnp.zeros((1, )) # pylint: disable=unused-variable
d = jnp.ones((1, )) # pylint: disable=unused-variable
e = jnp.zeros((1, 1)) # pylint: disable=unused-variable
f = jnp.ones((1, 1)) # pylint: disable=unused-variable
inferred = utils.pytype_infer_types(_PREAMBLE + code_saver.code)
self.assertEqual(inferred.a, 'Array0')
self.assertEqual(inferred.b, 'Array0')
self.assertEqual(inferred.c, 'Array1')
self.assertEqual(inferred.d, 'Array1')
self.assertEqual(inferred.e, 'Array2')
self.assertEqual(inferred.f, 'Array2')
def testMathUnaryOperator_ReturnCustomType(self):
with utils.SaveCodeAsString() as code_saver:
x: Tensor1[A1] = tf.zeros((1, ))
# Let's just test a representative subset.
a = tf.math.abs(x) # pylint: disable=unused-variable
b = tf.math.sin(x) # pylint: disable=unused-variable
c = tf.math.floor(x) # pylint: disable=unused-variable
d = tf.math.round(x) # pylint: disable=unused-variable
e = tf.math.sign(x) # pylint: disable=unused-variable
inferred = utils.pytype_infer_types(_PREAMBLE + code_saver.code)
expected = 'Tensor1[A1]'
self.assertEqual(inferred.a, expected)
self.assertEqual(inferred.b, expected)
self.assertEqual(inferred.c, expected)
self.assertEqual(inferred.d, expected)
self.assertEqual(inferred.e, expected)
def testUnaryOperator_ReturnCustomType(self):
with utils.SaveCodeAsString() as code_saver:
x: Array1[A1] = jnp.zeros((1, ))
# Let's just test a representative subset.
a = jnp.abs(x) # pylint: disable=unused-variable
b = jnp.sin(x) # pylint: disable=unused-variable
c = jnp.floor(x) # pylint: disable=unused-variable
d = jnp.ones_like(x) # pylint: disable=unused-variable
e = jnp.round(x) # pylint: disable=unused-variable
f = jnp.sign(x) # pylint: disable=unused-variable
inferred = utils.pytype_infer_types(_PREAMBLE + code_saver.code)
expected = 'Array1[A1]'
self.assertEqual(inferred.a, expected)
self.assertEqual(inferred.b, expected)
self.assertEqual(inferred.c, expected)
self.assertEqual(inferred.d, expected)
self.assertEqual(inferred.e, expected)
self.assertEqual(inferred.f, expected)
def testShapeAttribute_HasTypeTensorShape(self):
with utils.SaveCodeAsString() as code_saver:
x0 = tf.zeros(())
x1 = tf.zeros((1, ))
x2 = tf.zeros((1, 2))
x3 = tf.zeros((1, 2, 3))
x4 = tf.zeros((1, 2, 3, 4))
x5 = tf.zeros((1, 2, 3, 4, 5))
x0_shape = x0.shape # pylint: disable=unused-variable
x1_shape = x1.shape # pylint: disable=unused-variable
x2_shape = x2.shape # pylint: disable=unused-variable
x3_shape = x3.shape # pylint: disable=unused-variable
x4_shape = x4.shape # pylint: disable=unused-variable
x5_shape = x5.shape # pylint: disable=unused-variable
inferred = utils.pytype_infer_types(_PREAMBLE + code_saver.code)
self.assertEqual('tensorflow.TensorShape', inferred.x0_shape)
self.assertEqual('tensorflow.TensorShape', inferred.x1_shape)
self.assertEqual('tensorflow.TensorShape', inferred.x2_shape)
self.assertEqual('tensorflow.TensorShape', inferred.x3_shape)
self.assertEqual('tensorflow.TensorShape', inferred.x4_shape)
self.assertEqual('tensorflow.TensorShape', inferred.x5_shape)
