Repository: google/jax

Issue:
Inconsistent casting of floats with jit compilation.
```
def mul(R, c):
  return R * c

R = np.array([1, 2, 3], dtype=np.float32)
c = 0.1

print(mul(R, c).dtype) # Evaluates to np.float32 as expected.
print(jit(mul)(R, c).dtype) # Evaluates to np.float64.
```




Discussion:
Thanks for catching this.

I verified that this only happens with jax_enable_x64 switched on. I believe this is a Python scalars issue (which our tests don't cover).

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -735,7 +735,8 @@ def _get_min_identity(dtype):
     return onp.array(True, onp.bool_)
 
 def _reduce_sum(operand, axes):
-  return reduce_sum_p.bind(operand, axes=tuple(axes), input_shape=operand.shape)
+  return reduce_sum_p.bind(operand, axes=tuple(axes),
+                           input_shape=onp.shape(operand))
 
 def _reduce_prod(operand, axes):
   return reduce_prod_p.bind(operand, axes=tuple(axes))
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -150,17 +150,11 @@ def _promote_dtypes(*args):
   if len(args) < 2:
     return args
   else:
-    from_dtypes = (x if type(x) in _builtin_numeric_types else _dtype(x)
-                   for x in args)
+    from_dtypes = map(_dtype, args)
     to_dtype = xla_bridge.canonicalize_dtype(result_type(*from_dtypes))
     return [lax.convert_element_type(x, to_dtype)
             if _dtype(x) != to_dtype else x for x in args]
 
-if six.PY3:
-  _builtin_numeric_types = (int, float, complex)
-else:
-  _builtin_numeric_types = (int, float, long, complex)
-
 def _promote_to_result_dtype(op, *args):
   """Convenience function to promote args directly to the op's result dtype."""
   to_dtype = _result_dtype(op, *args)
@@ -1163,10 +1157,11 @@ def pad(array, pad_width, mode, constant_values=0):
     raise NotImplementedError(msg.format(mode))
 
   array = asarray(array)
-  pad_width = onp.broadcast_to(onp.asarray(pad_width), (array.ndim, 2))
-  constant_values = broadcast_to(asarray(constant_values), (array.ndim, 2))
-  for i in xrange(array.ndim):
-    widths = [(0, 0, 0)] * array.ndim
+  nd = ndim(array)
+  pad_width = onp.broadcast_to(onp.asarray(pad_width), (nd, 2))
+  constant_values = broadcast_to(asarray(constant_values), (nd, 2))
+  for i in xrange(nd):
+    widths = [(0, 0, 0)] * nd
     widths[i] = (pad_width[i, 0], 0, 0)
     array = lax.pad(array, constant_values[i, 0], widths)
     widths[i] = (0, pad_width[i, 1], 0)
@@ -1191,13 +1186,13 @@ def stack(arrays, axis=0):
 
 @_wraps(onp.tile)
 def tile(a, reps):
-    if isinstance(reps, int):
-        reps = (reps,)
-    a = a[(None,) * (len(reps) - a.ndim)]
-    reps = (1,) * (a.ndim - len(reps)) + reps
-    for i, rep in enumerate(reps):
-        a = concatenate([a] * rep, axis=i)
-    return a
+  if isinstance(reps, int):
+    reps = (reps,)
+  a = reshape(a, (1,) * (len(reps) - ndim(a)) + shape(a))
+  reps = (1,) * (ndim(a) - len(reps)) + reps
+  for i, rep in enumerate(reps):
+    a = concatenate([a] * rep, axis=i)
+  return a
 
 @_wraps(onp.concatenate)
 def concatenate(arrays, axis=0):
@@ -1338,9 +1333,7 @@ def array_equal(a1, a2):
     a1, a2 = asarray(a1), asarray(a2)
   except Exception:
     return False
-  if a1.shape != a2.shape:
-    return False
-  return asarray(a1==a2).all()
+  return shape(a1) == shape(a2) and all(asarray(a1 == a2))
 
 
 # We can't create uninitialized arrays in XLA; use zeros for empty.
@@ -1864,37 +1857,15 @@ def cross(a, b, axisa=-1, axisb=-1, axisc=-1, axis=None):
 
 @_wraps(onp.kron)
 def kron(a, b):
-  a_shape = shape(a)
-  b_shape = shape(b)
-  a_ndims = len(a_shape)
-  b_ndims = len(b_shape)
-  a = array(a)
-  b = array(b)
-  d = _min(a_ndims, b_ndims)
-  if d == 0:
-    return a * b
-  a_broadcast_dims = list(range(a_ndims - d, a_ndims + d, 2))
-  a_broadcast_shape = onp.ones(a_ndims + d, dtype=onp.int64)
-  a_broadcast_shape[:-2*d] = a_shape[:-d]
-  a_broadcast_shape[a_broadcast_dims] = a_shape[-d:]
-
-  b_broadcast_dims = list(range(b_ndims -d + 1, b_ndims + d + 1, 2))
-  b_broadcast_shape = onp.ones(b_ndims + d, dtype=onp.int64)
-  b_broadcast_shape[:-2*d] = b_shape[:-d]
-  b_broadcast_shape[b_broadcast_dims] = b_shape[-d:]
-
-  if a_ndims > b_ndims:
-    out_shape = onp.array(a_shape, dtype=onp.int64)
-    out_shape[-d:] *= onp.array(b_shape, dtype=onp.int64)
-  else:
-    out_shape = onp.array(b_shape, dtype=onp.int64)
-    out_shape[-d:] *= onp.array(a_shape, dtype=onp.int64)
-
-  a_broadcast = lax.broadcast_in_dim(
-    a, a_broadcast_shape, list(range(a_ndims - d)) + a_broadcast_dims)
-  b_broadcast = lax.broadcast_in_dim(
-    b, b_broadcast_shape, list(range(b_ndims - d)) + b_broadcast_dims)
-  return lax.reshape(a_broadcast * b_broadcast, out_shape)
+  a, b = _promote_dtypes(a, b)
+  if ndim(a) < ndim(b):
+    a = reshape(a, (1,) * (ndim(b) - ndim(a)) + shape(a))
+  elif ndim(b) < ndim(a):
+    b = reshape(b, (1,) * (ndim(a) - ndim(b)) + shape(b))
+  a_reshaped = reshape(a, [i for d in shape(a) for i in (d, 1)])
+  b_reshaped = reshape(b, [i for d in shape(b) for i in (1, d)])
+  out_shape = tuple(onp.multiply(shape(a), shape(b)))
+  return reshape(lax.mul(a_reshaped, b_reshaped), out_shape)
 
 
 @_wraps(onp.vander)
diff --git a/jax/scipy/special.py b/jax/scipy/special.py
--- a/jax/scipy/special.py
+++ b/jax/scipy/special.py
@@ -50,7 +50,7 @@ def expit(x):
   x = asarray(x)
   one = lax._const(x, 1)
   return lax.div(one, lax.add(one, lax.exp(lax.neg(x))))
-ad.defjvp2(expit.primitive, lambda g, ans, x: g * ans * (1 - ans))
+ad.defjvp2(expit.primitive, lambda g, ans, x: g * ans * (lax._const(ans, 1) - ans))
 batching.defvectorized(expit.primitive)
 
 


Test patch:
diff --git a/jax/test_util.py b/jax/test_util.py
--- a/jax/test_util.py
+++ b/jax/test_util.py
@@ -205,29 +205,39 @@ def format_test_name_suffix(opname, shapes, dtypes):
   return '{}_{}'.format(opname.capitalize(), '_'.join(arg_descriptions))
 
 
-class _NumpyScalar(object):
-
-  def __len__(self):
-    return 0
-
-# A special singleton "shape" that denotes numpy scalars. Numpy scalars are not
-# identical to 0-D arrays, and we want to write tests that exercise both paths.
+# We use special symbols, represented as singleton objects, to distinguish
+# between NumPy scalars, Python scalars, and 0-D arrays.
+class ScalarShape(object):
+  def __len__(self): return 0
+class _NumpyScalar(ScalarShape): pass
+class _PythonScalar(ScalarShape): pass
 NUMPY_SCALAR_SHAPE = _NumpyScalar()
+PYTHON_SCALAR_SHAPE = _PythonScalar()
 
 
 def _dims_of_shape(shape):
   """Converts `shape` to a tuple of dimensions."""
-  return shape if shape != NUMPY_SCALAR_SHAPE else ()
+  if type(shape) in (list, tuple):
+    return shape
+  elif isinstance(shape, ScalarShape):
+    return ()
+  else:
+    raise TypeError(type(shape))
 
 
 def _cast_to_shape(value, shape, dtype):
   """Casts `value` to the correct Python type for `shape` and `dtype`."""
-  if shape != NUMPY_SCALAR_SHAPE:
+  if shape is NUMPY_SCALAR_SHAPE:
+    # explicitly cast to NumPy scalar in case `value` is a Python scalar.
+    return dtype(value)
+  elif shape is PYTHON_SCALAR_SHAPE:
+    # explicitly cast to Python scalar via https://stackoverflow.com/a/11389998
+    return onp.asarray(value).item()
+  elif type(shape) in (list, tuple):
+    assert onp.shape(value) == tuple(shape)
     return value
   else:
-    # A numpy scalar was requested. Explicitly cast in case `value` is a Python
-    # scalar.
-    return dtype(value)
+    raise TypeError(type(shape))
 
 
 def dtype_str(dtype):
@@ -235,14 +245,17 @@ def dtype_str(dtype):
 
 
 def format_shape_dtype_string(shape, dtype):
-  if shape == NUMPY_SCALAR_SHAPE:
+  if shape is NUMPY_SCALAR_SHAPE:
     return dtype_str(dtype)
-
-  if onp.isscalar(shape):
-    shapestr = str(shape) + ','
-  else:
+  elif shape is PYTHON_SCALAR_SHAPE:
+    return 'py' + dtype_str(dtype)
+  elif type(shape) in (list, tuple):
     shapestr = ','.join(str(dim) for dim in shape)
-  return '{}[{}]'.format(dtype_str(dtype), shapestr)
+    return '{}[{}]'.format(dtype_str(dtype), shapestr)
+  elif type(shape) is int:
+    return '{}[{},]'.format(dtype_str(dtype), shape)
+  else:
+    raise TypeError(type(shape))
 
 
 def _rand_dtype(rand, shape, dtype, scale=1., post=lambda x: x):
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -20,6 +20,7 @@
 import functools
 from functools import partial
 import itertools
+import operator
 import unittest
 from unittest import SkipTest
 
@@ -43,7 +44,7 @@
 nonempty_array_shapes = [()] + nonempty_nonscalar_array_shapes
 empty_array_shapes = [(0,), (0, 4), (3, 0),]
 
-scalar_shapes = [jtu.NUMPY_SCALAR_SHAPE]
+scalar_shapes = [jtu.NUMPY_SCALAR_SHAPE, jtu.PYTHON_SCALAR_SHAPE]
 array_shapes = nonempty_array_shapes + empty_array_shapes
 nonzerodim_shapes = nonempty_nonscalar_array_shapes + empty_array_shapes
 nonempty_shapes = scalar_shapes + nonempty_array_shapes
@@ -201,40 +202,43 @@ def op_record(name, nargs, dtypes, shapes, rng, diff_modes, test_name=None,
 
 JAX_OPERATOR_OVERLOADS = [
     op_record("__add__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
-    op_record("__radd__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
     op_record("__sub__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
-    op_record("__rsub__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
     op_record("__mul__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
-    op_record("__rmul__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
     op_record("__eq__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
     op_record("__ne__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
-    op_record("__lt__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
-    op_record("__gt__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
-    op_record("__ge__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
+    op_record("__lt__", 2, default_dtypes, all_shapes, jtu.rand_default(), []),
+    op_record("__gt__", 2, default_dtypes, all_shapes, jtu.rand_default(), []),
+    op_record("__ge__", 2, default_dtypes, all_shapes, jtu.rand_default(), []),
     op_record("__neg__", 1, number_dtypes, all_shapes, jtu.rand_default(), []),
     op_record("__pow__", 2, inexact_dtypes, all_shapes, jtu.rand_positive(), []),
-    op_record("__rpow__", 2, inexact_dtypes, all_shapes, jtu.rand_positive(), []),
     op_record("__mod__", 2, default_dtypes, all_shapes, jtu.rand_nonzero(), []),
-    op_record("__rmod__", 2, default_dtypes, all_shapes, jtu.rand_nonzero(), []),
-    op_record("__floordiv__", 2, number_dtypes, all_shapes, jtu.rand_nonzero(), []),
-    op_record("__rfloordiv__", 2, number_dtypes, all_shapes, jtu.rand_nonzero(), []),
+    op_record("__floordiv__", 2, default_dtypes, all_shapes, jtu.rand_nonzero(), []),
     op_record("__truediv__", 2, number_dtypes, all_shapes, jtu.rand_nonzero(), []),
-    op_record("__rtruediv__", 2, number_dtypes, all_shapes, jtu.rand_nonzero(), []),
     op_record("__abs__", 1, number_dtypes, all_shapes, jtu.rand_default(), []),
     # TODO(mattjj): __invert__ fails on bool dtypes because ~True == -2
     op_record("__invert__", 1, int_dtypes, all_shapes, jtu.rand_default(), []),
     # TODO(mattjj): investigate these failures
     # op_record("__or__", 2, number_dtypes, all_shapes, jtu.rand_bool(), []),
-    # op_record("__ror__", 2, number_dtypes, all_shapes, jtu.rand_bool(), []),
     # op_record("__and__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
-    # op_record("__rand__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
     # op_record("__xor__", 2, number_dtypes, all_shapes, jtu.rand_bool(), []),
-    # op_record("__rxor__", 2, number_dtypes, all_shapes, jtu.rand_bool(), []),
     # op_record("__divmod__", 2, number_dtypes, all_shapes, jtu.rand_nonzero(), []),
-    # op_record("__rdivmod__", 2, number_dtypes, all_shapes, jtu.rand_nonzero(), []),
     # TODO(mattjj): lshift, rshift
 ]
 
+JAX_RIGHT_OPERATOR_OVERLOADS = [
+    op_record("__radd__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
+    op_record("__rsub__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
+    op_record("__rmul__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
+    op_record("__rpow__", 2, inexact_dtypes, all_shapes, jtu.rand_positive(), []),
+    op_record("__rmod__", 2, default_dtypes, all_shapes, jtu.rand_nonzero(), []),
+    op_record("__rfloordiv__", 2, default_dtypes, all_shapes, jtu.rand_nonzero(), []),
+    op_record("__rtruediv__", 2, number_dtypes, all_shapes, jtu.rand_nonzero(), []),
+    # op_record("__ror__", 2, number_dtypes, all_shapes, jtu.rand_bool(), []),
+    # op_record("__rand__", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
+    # op_record("__rxor__", 2, number_dtypes, all_shapes, jtu.rand_bool(), []),
+    # op_record("__rdivmod__", 2, number_dtypes, all_shapes, jtu.rand_nonzero(), []),
+]
+
 numpy_version = tuple(map(int, onp.version.version.split('.')))
 if numpy_version >= (1, 15):
   JAX_COMPOUND_OP_RECORDS += [
@@ -249,6 +253,8 @@ def op_record(name, nargs, dtypes, shapes, rng, diff_modes, test_name=None,
 if six.PY2:
   JAX_OPERATOR_OVERLOADS += [
     op_record("__div__", 2, number_dtypes, all_shapes, jtu.rand_nonzero(), []),
+  ]
+  JAX_RIGHT_OPERATOR_OVERLOADS += [
     op_record("__rdiv__", 2, number_dtypes, all_shapes, jtu.rand_nonzero(), []),
   ]
 
@@ -293,7 +299,7 @@ def _GetArgsMaker(self, rng, shapes, dtypes):
                                                       dtypes),
          "rng": rec.rng, "shapes": shapes, "dtypes": dtypes,
          "onp_op": getattr(onp, rec.name), "lnp_op": getattr(lnp, rec.name),
-          "check_dtypes": rec.check_dtypes}
+         "check_dtypes": rec.check_dtypes}
         for shapes in filter(
           _shapes_are_broadcast_compatible,
           CombosWithReplacement(rec.shapes, rec.nargs))
@@ -302,8 +308,9 @@ def _GetArgsMaker(self, rng, shapes, dtypes):
                                  JAX_COMPOUND_OP_RECORDS)))
   def testOp(self, onp_op, lnp_op, rng, shapes, dtypes, check_dtypes):
     args_maker = self._GetArgsMaker(rng, shapes, dtypes)
+    py_scalar_arg = jtu.PYTHON_SCALAR_SHAPE in shapes
     self._CheckAgainstNumpy(onp_op, lnp_op, args_maker,
-                            check_dtypes=check_dtypes)
+                            check_dtypes=check_dtypes and not py_scalar_arg)
     self._CompileAndCheck(lnp_op, args_maker, check_dtypes=check_dtypes)
 
   @parameterized.named_parameters(itertools.chain.from_iterable(
@@ -318,8 +325,27 @@ def testOp(self, onp_op, lnp_op, rng, shapes, dtypes, check_dtypes):
       for rec in JAX_OPERATOR_OVERLOADS))
   def testOperatorOverload(self, name, rng, shapes, dtypes):
     args_maker = self._GetArgsMaker(rng, shapes, dtypes)
-    fun = lambda x, *xs: getattr(x, name)(*xs)
-    self._CompileAndCheck(fun, args_maker, check_dtypes=True)
+    fun = lambda *xs: getattr(operator, name.strip('_'))(*xs)
+    self._CompileAndCheck(fun, args_maker,
+                          check_dtypes=jtu.PYTHON_SCALAR_SHAPE not in shapes)
+
+  @parameterized.named_parameters(itertools.chain.from_iterable(
+      jtu.cases_from_list(
+        {"testcase_name": jtu.format_test_name_suffix(rec.test_name, shapes,
+                                                      dtypes),
+         "rng": rec.rng, "shapes": shapes, "dtypes": dtypes, "name": rec.name}
+        for shapes in filter(
+          _shapes_are_broadcast_compatible,
+          CombosWithReplacement(rec.shapes, rec.nargs))
+        for dtypes in CombosWithReplacement(rec.dtypes, rec.nargs))
+      for rec in JAX_RIGHT_OPERATOR_OVERLOADS))
+  def testRightOperatorOverload(self, name, rng, shapes, dtypes):
+    if shapes[1] is jtu.PYTHON_SCALAR_SHAPE:
+      raise SkipTest()  # TODO(mattjj): clean up
+    args_maker = self._GetArgsMaker(rng, shapes, dtypes)
+    fun = lambda fst, snd: getattr(snd, name)(fst)
+    self._CompileAndCheck(fun, args_maker,
+                          check_dtypes=jtu.PYTHON_SCALAR_SHAPE not in shapes)
 
   @parameterized.named_parameters(itertools.chain.from_iterable(
       jtu.cases_from_list(
@@ -339,7 +365,8 @@ def testBitwiseOp(self, onp_op, lnp_op, rng, shapes, dtypes):
         onp.iinfo(dtype).bits == 64 for dtype in dtypes):
       self.skipTest("x64 types are disabled by jax_enable_x64")
     args_maker = self._GetArgsMaker(rng, shapes, dtypes)
-    self._CheckAgainstNumpy(onp_op, lnp_op, args_maker, check_dtypes=True)
+    self._CheckAgainstNumpy(onp_op, lnp_op, args_maker,
+                            check_dtypes=jtu.PYTHON_SCALAR_SHAPE not in shapes)
     self._CompileAndCheck(lnp_op, args_maker, check_dtypes=True)
 
   @parameterized.named_parameters(jtu.cases_from_list(
@@ -1405,6 +1432,10 @@ def testIssue330(self):
     self.assertEqual(x[0, 0], 1)
 
   def testScalarDtypePromotion(self):
+    # disabled this test after https://github.com/google/jax/issues/732
+    msg = ("jax.numpy differs from numpy in promotion rules for Python scalars."
+           " See https://github.com/google/jax/issues/732.")
+    raise SkipTest(msg)
     orig_numpy_result = (1 + onp.eye(1, dtype=onp.float32)).dtype
     jax_numpy_result = (1 + lnp.eye(1, dtype=lnp.float32)).dtype
     self.assertEqual(orig_numpy_result, jax_numpy_result)
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1364,7 +1364,10 @@ def testFilledConstant(self, shape, fill_value, dtype):
           jtu.format_shape_dtype_string(shape, dtype), dimension),
        "shape": shape, "dtype": dtype, "dimension": dimension}
       for dtype in default_dtypes
-      for shape in [(), (3,), (2, 3), (2, 3, 4), (1001, 1001), (101, 101, 101)]
+      for shape in [(), (3,), (2, 3), (2, 3, 4),
+                    # TODO(mattjj): re-enable
+                    # (1001, 1001), (101, 101, 101),
+                    ]
       for dimension in range(len(shape))))
   def testIotaConstant(self, dtype, shape, dimension):
     make_const = lambda: lax.broadcasted_iota(dtype, shape, dimension)