Not compatible with np.tensordot

Environments:

• jaxlib: 0.1.61+cuda110
• jax: 0.2.10
• dm-haiku: 0.0.4.dev0
• numpy: 1.19.5

To reproduce the bug:

• code is given as bellow
• you might need to run it multiple times

import jax
import haiku as hk
import numpy as np

def net(x):
  torso_net = hk.Conv2D(32, kernel_shape=[8, 8], stride=[4, 4], padding='VALID')
  return  hk.BatchApply(torso_net)(x)

def main():
  init_fn, apply_fn = hk.without_apply_rng(
    hk.transform(lambda obs: net(obs))
  )

  # Uncomment can also reproduce the bug
  # @jax.jit
  def initial_params(rng_key):
    _ = np.tensordot(np.zeros((84,84,3)), [0., 0., 0.], (-1, 0))
    return init_fn(rng_key, np.zeros((1,1,84,84,4)))

  rng_key = jax.random.PRNGKey(428)
  params = initial_params(rng_key)

if __name__ == '__main__':
  main()

The full error messages/tracebacks.

• with @jax.jit

Segmentation fault (core dumped)

• without @jax.jit

/usr/local/lib/python3.6/dist-packages/jax/_src/numpy/lax_numpy.py:2970: UserWarning: Explicitly requested dtype float64 requested in zeros is not available, and will be truncated to dtype float32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/google/jax#current-gotchas for more.
  lax._check_user_dtype_supported(dtype, "zeros")
Segmentation fault (core dumped)

[bingykang]

Looks like the bug is gone.

[mattjj]

Can you try this again on the latest jaxlib==0.1.64? I think you can just write pip install --upgrade jaxlib==0.1.64+cuda110 -f https://storage.googleapis.com/jax-releases/jax_releases.html to use the same cuda110 version you have now.