Python 3 compatibility issues

[cclauss]

flake8 testing of https://github.com/google/jax on Python 3.7.1

$ flake8 . --count --select=E901,E999,F821,F822,F823 --show-source --statistics

./jax/lax.py:481:25: F821 undefined name '_ndim'
  start_indices = [0] * _ndim(operand)
                        ^
./jax/lax.py:487:6: F821 undefined name '_ndim'
  if _ndim(update) != _ndim(operand):
     ^
./jax/lax.py:487:23: F821 undefined name '_ndim'
  if _ndim(update) != _ndim(operand):
                      ^
./jax/lax.py:488:12: F821 undefined name '_ndim'
    assert _ndim(update) + 1 == _ndim(operand)
           ^
./jax/lax.py:488:33: F821 undefined name '_ndim'
    assert _ndim(update) + 1 == _ndim(operand)
                                ^
./jax/lax.py:489:17: F821 undefined name '_ndim'
    ax = axis % _ndim(operand)
                ^
./jax/lax.py:1955:35: F821 undefined name 'c'
  select = _reduction_computation(c, select_jaxpr, select_consts, init_value)
                                  ^
./jax/lax.py:1956:36: F821 undefined name 'c'
  scatter = _reduction_computation(c, scatter_jaxpr, scatter_consts, init_value)
                                   ^
./jax/lax.py:1957:10: F821 undefined name 'c'
  return c.SelectAndScatter(operand, select, window_dimensions, window_strides,
         ^
./jax/lax.py:2156:32: F821 undefined name 'name'
    raise TypeError(msg.format(name, len(lhs_shape), len(rhs_shape)))
                               ^
./jax/abstract_arrays.py:59:46: F821 undefined name 'long'
    _long    = concretization_function_error(long)
                                             ^
./jax/core.py:163:7: F821 undefined name 'print_trace_stack'
      print_trace_stack()
      ^
./jax/core.py:371:7: F821 undefined name 'print_trace_stack'
      print_trace_stack()
./jax/interpreters/xla.py:252:48: F821 undefined name 'long'
    __long__ = partialmethod(forward_to_value, long)
                                               ^
./jax/interpreters/ad.py:189:20: F821 undefined name 'JaxTuple'
        return xt, JaxTuple(map(zeros_like_jaxval, xt))
                   ^
./jax/interpreters/ad.py:196:16: F821 undefined name 'JaxTuple'
        return JaxTuple(map(zeros_like_jaxval, yt)), yt
               ^
./jax/numpy/lax_numpy.py:379:41: F821 undefined name 'isfortran'
    dims = onp.arange(ndim(a))[::-1] if isfortran(a) else onp.arange(ndim(a))
                                        ^
./examples/mnist_vae.py:113:21: E999 SyntaxError: invalid syntax
    def body_fun(i, (rng, opt_state, images)):
                    ^
./examples/resnet50.py:124:12: F821 undefined name 'xrange'
  for i in xrange(num_steps):
           ^
./tests/minmax_test.py:46:16: F821 undefined name 'fax'
    infeeder = fax.make_infeed_from_sequence(
               ^
1     E999 SyntaxError: invalid syntax
19    F821 undefined name 'xrange'
20

[mattjj]

Thanks for raising this, and for starting the PRs #2, #3, and #4!

Those are all fixes we need to make, but it might take us some time to merge the PRs because we're still ironing out our code syncing process. In other words, we're not quite ready for PRs yet, but we'll try to merge these if and when we can.

Thanks for your interest and your work to improve JAX (even before we're ready for it!).

[cclauss]

No rush... Take the time that you need to get it right.